Additional methods and fixes to DataFrameMixin

mesa_frames/abstract/mixin.py

@@ -1,11 +1,13 @@
 from abc import ABC, abstractmethod
+from collections.abc import Collection, Iterator, Sequence
 from copy import copy, deepcopy
-
-from typing_extensions import Any, Self
 from typing import Literal
-from collections.abc import Collection, Iterator, Sequence
 
-from mesa_frames.types_ import BoolSeries, DataFrame, Mask, Series
+from typing_extensions import Any, Self, overload
+
+from collections.abc import Hashable
+
+from mesa_frames.types_ import BoolSeries, DataFrame, Index, Mask, Series
 
 
 class CopyMixin(ABC):
@@ -149,38 +151,119 @@ def __deepcopy__(self, memo: dict) -> Self:
 
 
 class DataFrameMixin(ABC):
+    def _df_remove(self, df: DataFrame, mask: Mask, index_cols: str) -> DataFrame:
+        return self._df_get_masked_df(df, index_cols, mask, negate=True)
+
     @abstractmethod
-    def _df_add_columns(
-        self, original_df: DataFrame, new_columns: list[str], data: Any
+    def _df_add(
+        self,
+        df: DataFrame,
+        other: DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_all(
+        self,
+        df: DataFrame,
+        name: str,
+        axis: str = "columns",
+        index_cols: str | list[str] | None = None,
     ) -> DataFrame: ...
 
+    @abstractmethod
+    def _df_column_names(self, df: DataFrame) -> list[str]: ...
+
     @abstractmethod
     def _df_combine_first(
-        self, original_df: DataFrame, new_df: DataFrame, index_cols: list[str]
+        self, original_df: DataFrame, new_df: DataFrame, index_cols: str | list[str]
     ) -> DataFrame: ...
 
+    @overload
     @abstractmethod
     def _df_concat(
         self,
-        dfs: Collection[DataFrame],
+        objs: Collection[Series],
         how: Literal["horizontal"] | Literal["vertical"] = "vertical",
         ignore_index: bool = False,
+        index_cols: str | None = None,
+    ) -> Series: ...
+
+    @overload
+    @abstractmethod
+    def _df_concat(
+        self,
+        objs: Collection[DataFrame],
+        how: Literal["horizontal"] | Literal["vertical"] = "vertical",
+        ignore_index: bool = False,
+        index_cols: str | None = None,
     ) -> DataFrame: ...
 
+    @abstractmethod
+    def _df_concat(
+        self,
+        objs: Collection[DataFrame] | Collection[Series],
+        how: Literal["horizontal"] | Literal["vertical"] = "vertical",
+        ignore_index: bool = False,
+        index_cols: str | None = None,
+    ) -> DataFrame | Series: ...
+
+    @abstractmethod
+    def _df_contains(
+        self,
+        df: DataFrame,
+        column: str,
+        values: Sequence[Any],
+    ) -> BoolSeries: ...
+
     @abstractmethod
     def _df_constructor(
         self,
         data: Sequence[Sequence] | dict[str | Any] | None = None,
         columns: list[str] | None = None,
-        index_col: str | list[str] | None = None,
+        index: Index | None = None,
+        index_cols: str | list[str] | None = None,
         dtypes: dict[str, Any] | None = None,
     ) -> DataFrame: ...
 
+    @abstractmethod
+    def _df_div(
+        self,
+        df: DataFrame,
+        other: DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_drop_columns(
+        self,
+        df: DataFrame,
+        columns: str | list[str],
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_drop_duplicates(
+        self,
+        df: DataFrame,
+        subset: str | list[str] | None = None,
+        keep: Literal["first", "last", False] = "first",
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_filter(
+        self,
+        df: DataFrame,
+        condition: BoolSeries,
+        all: bool = True,
+    ) -> DataFrame: ...
+
     @abstractmethod
     def _df_get_bool_mask(
         self,
         df: DataFrame,
-        index_col: str,
+        index_cols: str | list[str],
         mask: Mask | None = None,
         negate: bool = False,
     ) -> BoolSeries: ...
@@ -189,21 +272,88 @@ def _df_get_bool_mask(
     def _df_get_masked_df(
         self,
         df: DataFrame,
-        index_col: str,
+        index_cols: str,
         mask: Mask | None = None,
-        columns: list[str] | None = None,
+        columns: str | list[str] | None = None,
         negate: bool = False,
     ) -> DataFrame: ...
 
+    @abstractmethod
+    def _df_groupby_cumcount(
+        self,
+        df: DataFrame,
+        by: str | list[str],
+    ) -> Series: ...
+
     @abstractmethod
     def _df_iterator(self, df: DataFrame) -> Iterator[dict[str, Any]]: ...
 
     @abstractmethod
-    def _df_norm(self, df: DataFrame) -> DataFrame: ...
+    def _df_join(
+        self,
+        left: DataFrame,
+        right: DataFrame,
+        index_cols: str | list[str] | None = None,
+        on: str | list[str] | None = None,
+        left_on: str | list[str] | None = None,
+        right_on: str | list[str] | None = None,
+        how: Literal["left"]
+        | Literal["right"]
+        | Literal["inner"]
+        | Literal["outer"]
+        | Literal["cross"] = "left",
+        suffix="_right",
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_mul(
+        self,
+        df: DataFrame,
+        other: DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    @overload
+    def _df_norm(
+        self,
+        df: DataFrame,
+        srs_name: str = "norm",
+        include_cols: Literal[False] = False,
+    ) -> Series: ...
+
+    @abstractmethod
+    @overload
+    def _df_norm(
+        self,
+        df: DataFrame,
+        srs_name: str = "norm",
+        include_cols: Literal[True] = False,
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_norm(
+        self,
+        df: DataFrame,
+        srs_name: str = "norm",
+        include_cols: bool = False,
+    ) -> Series | DataFrame: ...
 
     @abstractmethod
-    def _df_remove(
-        self, df: DataFrame, ids: Sequence[Any], index_col: str | None = None
+    def _df_rename_columns(
+        self,
+        df: DataFrame,
+        old_columns: list[str],
+        new_columns: list[str],
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_reset_index(
+        self,
+        df: DataFrame,
+        index_cols: str | list[str] | None = None,
+        drop: bool = False,
     ) -> DataFrame: ...
 
     @abstractmethod
@@ -217,6 +367,27 @@ def _df_sample(
         seed: int | None = None,
     ) -> DataFrame: ...
 
+    @abstractmethod
+    def _df_set_index(
+        self,
+        df: DataFrame,
+        index_name: str,
+        new_index: Sequence[Hashable] | None = None,
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_with_columns(
+        self,
+        original_df: DataFrame,
+        data: DataFrame
+        | Series
+        | Sequence[Sequence]
+        | dict[str | Any]
+        | Sequence[Any]
+        | Any,
+        new_columns: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
     @abstractmethod
     def _srs_constructor(
         self,
@@ -225,3 +396,16 @@ def _srs_constructor(
         dtype: Any | None = None,
         index: Sequence[Any] | None = None,
     ) -> Series: ...
+
+    @abstractmethod
+    def _srs_contains(
+        self,
+        srs: Sequence[Any],
+        values: Any | Sequence[Any],
+    ) -> BoolSeries: ...
+
+    @abstractmethod
+    def _srs_range(self, name: str, start: int, end: int, step: int = 1) -> Series: ...
+
+    @abstractmethod
+    def _srs_to_df(self, srs: Series, index: Index | None = None) -> DataFrame: ...