1
\$\begingroup\$

Below is my solution for managing multiple DataFrames, in an abstract enough way that it may apply to objects outside of a pandas.DataFrame hence the Protocol SupportsGetAndSetItemSequence.

This solution supports a MultiIndex index and assumes columns will be a single Index

Typically I would stuff all of the data into a single DataFrame but working with geospatial timeseries data often leads to a misaligned indices.

Note: I am using pyright for static type analysis and python 3.10

from __future__ import annotations

import enum
from typing import TypeVar, Generic, Any, overload, Sequence, TypeGuard, Protocol

from typing_extensions import Self
import pandas as pd
from pandas.core.indexes.category import CategoricalIndex


Key_T_co = TypeVar("Key_T_co", contravariant=True)


class SupportsGetAndSetItemSequence(Protocol[Key_T_co]):
    @overload
    def __getitem__(self, __key: list[Key_T_co]) -> Self:
        ...

    @overload
    def __getitem__(self, __key: Key_T_co) -> Self:
        ...

    def __setitem__(self, __key: list[Key_T_co] | Key_T_co, __value: Any) -> None:
        ...


class Coordinate(enum.Enum):
    @classmethod
    def to_index(cls) -> CategoricalIndex:
        return CategoricalIndex(cls.to_list(), name=cls, dtype="category")

    @classmethod
    def to_list(cls) -> list[Self]:
        return list(cls)


Coord_T = TypeVar("Coord_T", bound=Coordinate)
Value_T = TypeVar("Value_T", bound=SupportsGetAndSetItemSequence)



def is_coordinate_subclass(cls: type[Coord_T] | Any) -> TypeGuard[type[Coord_T]]:
    if cls is Coordinate:
        raise TypeError("Coordinate is not a valid Coordinate subclass")
    return isinstance(cls, type) and issubclass(cls, Coordinate)


NodeStore = dict[type[Coord_T], Value_T]


class Dataset(Generic[Coord_T, Value_T]):
    __slots__ = ("_nodes",)

    def __init__(self, nodes: NodeStore[Coord_T, Value_T] | None = None) -> None:
        self._nodes = nodes or {}

    @property
    def nodes(self) -> NodeStore[Coord_T, Value_T]:
        return self._nodes

    @staticmethod
    def _get_coordinate(coord: Any) -> type[Coord_T]:
        if is_coordinate_subclass(coord):
            return coord
        elif isinstance(coord, Sequence):
            return type(coord[0])
        else:
            return type(coord)

    def __getitem__(self, key: type[Coord_T] | Sequence[Coord_T] | Coord_T) -> Value_T:
        coord = self._get_coordinate(key)
        return self._nodes[coord] if coord == key else self._nodes[coord][key]

    def __setitem__(self, key: type[Coord_T] | Sequence[Coord_T] | Coord_T, value: Any) -> None:
        coord = self._get_coordinate(key)
        if coord == key or coord not in self.nodes.keys():
            self.nodes[coord] = value
        else:
            self.nodes[coord][key] = value

    def __repr__(self) -> str:
        return "\n".join(f"{self.__class__.__name__}[{k.__name__}]\n{v}" for k, v in self.nodes.items())


Usage

class ExampleA(Coordinate):
    one = "cat"
    two = "dog"
    three = "bird"


class ExampleB(Coordinate):
    one = "dog"


class ExampleC(Coordinate):
    one = "bird"


def main():
    ds = Dataset[ExampleA | ExampleB, pd.DataFrame](
        {
            ExampleA: pd.DataFrame(
                {ExampleA.one: [1]},
                index=pd.MultiIndex.from_tuples([(1, 2)], names=["x", "y"]),
            )
        }
    )

    ds[ExampleA].columns.rename(ExampleA, inplace=True)  # Ok
    ds[ExampleA.two] = [5]  # Ok
    ds[ExampleA.two] += ds[ExampleA.two]  # Ok
    ds[ExampleA.three] = ds[[ExampleA.one, ExampleA.two]].sum(axis=1)  # Ok
    
    ds[ExampleB.one] = pd.DataFrame({ExampleB.one: [2]}, columns=ExampleB.to_index())  # Ok
    
    df1 = ds[ExampleA]  # Ok
    df2 = ds[ExampleB]  # Ok
    print(ds)
    # Dataset[ExampleA]
    # <enum 'ExampleA'>  ExampleA.one  ExampleA.two  ExampleA.three
    # x y                                                          
    # 1 2                           1            10              11
    # Dataset[ExampleB]
    # <enum 'ExampleB'>  ExampleB.one
    # 0                             2
    ds[ExampleC] = pd.DataFrame({ExampleC.one: [3]})  # Argument of Type "ExampleC" cannot be assigned ...


if __name__ == "__main__":
    main()

\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.