microdf-python 1.0.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

microdf/microdataframe.py

@@ -12,6 +12,80 @@ from microdf.microseries import MicroSeries, MicroSeriesGroupBy
 logger = logging.getLogger(__name__)
 
 
+class _MicroLocIndexer:
+    """Custom loc indexer that returns MicroDataFrame with proper weights."""
+
+    def __init__(self, mdf: "MicroDataFrame"):
+        self._mdf = mdf
+        # Get the parent's loc indexer
+        self._parent_loc = pd.DataFrame.loc.fget(mdf)
+
+    def __getitem__(self, key):
+        # Use the parent DataFrame's loc indexer
+        result = self._parent_loc[key]
+
+        if isinstance(result, pd.DataFrame):
+            # Get the filtered weights based on the result's index
+            new_weights = self._mdf.weights.reindex(result.index)
+            return MicroDataFrame(result, weights=new_weights)
+        elif isinstance(result, pd.Series):
+            # Single row or column selected
+            if result.name in self._mdf.columns:
+                # Column was selected - return MicroSeries with all weights
+                return MicroSeries(result, weights=self._mdf.weights)
+            else:
+                # Row was selected - return as-is (scalar values for each col)
+                return result
+        else:
+            # Scalar value
+            return result
+
+    def __setitem__(self, key, value):
+        self._parent_loc[key] = value
+        self._mdf._link_all_weights()
+
+
+class _MicroILocIndexer:
+    """Custom iloc indexer that returns MicroDataFrame with proper weights."""
+
+    def __init__(self, mdf: "MicroDataFrame"):
+        self._mdf = mdf
+        # Get the parent's iloc indexer
+        self._parent_iloc = pd.DataFrame.iloc.fget(mdf)
+
+    def __getitem__(self, key):
+        # Use the parent DataFrame's iloc indexer
+        result = self._parent_iloc[key]
+
+        if isinstance(result, pd.DataFrame):
+            # Get the filtered weights based on the result's index
+            new_weights = self._mdf.weights.iloc[
+                self._mdf.index.get_indexer(result.index)
+            ]
+            new_weights = pd.Series(new_weights.values, index=result.index)
+            return MicroDataFrame(result, weights=new_weights)
+        elif isinstance(result, pd.Series):
+            # Single row or column selected
+            if isinstance(key, tuple) and len(key) == 2:
+                # df.iloc[:, col_idx] - column selection
+                row_key = key[0]
+                if isinstance(row_key, slice) and row_key == slice(None):
+                    # All rows selected for a column
+                    return MicroSeries(result, weights=self._mdf.weights)
+            # Check if this is a column (result index matches mdf index)
+            if result.index.equals(self._mdf.index):
+                return MicroSeries(result, weights=self._mdf.weights)
+            # Row selection - return as-is
+            return result
+        else:
+            # Scalar value
+            return result
+
+    def __setitem__(self, key, value):
+        self._parent_iloc[key] = value
+        self._mdf._link_all_weights()
+
+
 class MicroDataFrame(pd.DataFrame):
     def __init__(self, *args, weights=None, **kwargs):
         """A DataFrame-inheriting class for weighted microdata. Weights can be
@@ -26,6 +100,23 @@ class MicroDataFrame(pd.DataFrame):
         self._link_all_weights()
         self.override_df_functions()
 
+    @property
+    def loc(self) -> _MicroLocIndexer:
+        """Label-based indexer that preserves MicroDataFrame type and weights.
+
+        :return: Custom loc indexer for MicroDataFrame
+        """
+        return _MicroLocIndexer(self)
+
+    @property
+    def iloc(self) -> _MicroILocIndexer:
+        """Integer-based indexer that preserves MicroDataFrame type and
+        weights.
+
+        :return: Custom iloc indexer for MicroDataFrame
+        """
+        return _MicroILocIndexer(self)
+
     def override_df_functions(self) -> None:
         """Override DataFrame functions to work with weighted operations."""
         for name in MicroSeries.FUNCTIONS:
@@ -643,6 +734,7 @@ class MicroDataFrame(pd.DataFrame):
 
 class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy):
     def _init(self, by: Union[str, List]):
+        self._by = by
         self.columns = list(self.obj.columns)
         if isinstance(by, list):
             for column in by:
@@ -656,6 +748,10 @@ class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy):
             for col in self.columns
             if pd.api.types.is_numeric_dtype(self.obj[col])
         ]
+        # Store reference to weights groupby for column selection
+        self._weights_groupby = copy.deepcopy(
+            super().__getitem__("__tmp_weights")
+        )
         for fn_name in MicroSeries.SCALAR_FUNCTIONS:
 
             def get_fn(name):
@@ -669,11 +765,9 @@ class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy):
                         except Exception:
                             # Skip columns that can't be aggregated
                             pass
-                    return (
-                        MicroDataFrame(results)
-                        if results
-                        else MicroDataFrame()
-                    )
+                    # Return plain DataFrame - aggregated results don't have
+                    # per-row weights (weights were already applied)
+                    return pd.DataFrame(results) if results else pd.DataFrame()
 
                 return fn
 
@@ -691,12 +785,99 @@ class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy):
                         except Exception:
                             # Skip columns that can't be aggregated
                             pass
-                    return (
-                        MicroDataFrame(results)
-                        if results
-                        else MicroDataFrame()
-                    )
+                    # Return plain DataFrame - aggregated results don't have
+                    # per-row weights (weights were already applied)
+                    return pd.DataFrame(results) if results else pd.DataFrame()
 
                 return fn
 
             setattr(self, fn_name, get_fn(fn_name))
+
+    def __getitem__(
+        self, key: Union[str, List]
+    ) -> Union["MicroSeriesGroupBy", "MicroDataFrameGroupBy"]:
+        """Select columns from the groupby object while preserving weights.
+
+        This ensures that operations like groupby(col)["y"].sum() or
+        groupby(col)[["y"]].sum() use weighted aggregation.
+
+        :param key: Column name or list of column names
+        :return: MicroSeriesGroupBy for single column, MicroDataFrameGroupBy
+            for multiple columns
+        """
+        if isinstance(key, str):
+            # Single column - return MicroSeriesGroupBy
+            result = super().__getitem__(key)
+            result.__class__ = MicroSeriesGroupBy
+            result._init()
+            result.weights = self._weights_groupby
+            return result
+        else:
+            # Multiple columns - return a new MicroDataFrameGroupBy
+            # with only the selected columns
+            result = super().__getitem__(key)
+            result.__class__ = MicroDataFrameGroupBy
+            # Re-initialize with the subset of columns
+            result._by = self._by
+            result.columns = list(key) if hasattr(key, "__iter__") else [key]
+            result.numeric_columns = [
+                col
+                for col in result.columns
+                if pd.api.types.is_numeric_dtype(result.obj[col])
+            ]
+            result._weights_groupby = self._weights_groupby
+            # Set up the column attributes as MicroSeriesGroupBy
+            for col in result.columns:
+                col_gb = super().__getitem__(col)
+                col_gb.__class__ = MicroSeriesGroupBy
+                col_gb._init()
+                col_gb.weights = self._weights_groupby
+                setattr(result, col, col_gb)
+            # Set up the scalar and vector functions
+            for fn_name in MicroSeries.SCALAR_FUNCTIONS:
+
+                def get_scalar_fn(name, res):
+                    def fn(*args, **kwargs):
+                        results = {}
+                        for col in res.numeric_columns:
+                            try:
+                                results[col] = getattr(
+                                    getattr(res, col), name
+                                )(*args, **kwargs)
+                            except Exception:
+                                pass
+                        # Return plain DataFrame - aggregated results don't
+                        # have per-row weights (weights were already applied)
+                        return (
+                            pd.DataFrame(results)
+                            if results
+                            else pd.DataFrame()
+                        )
+
+                    return fn
+
+                setattr(result, fn_name, get_scalar_fn(fn_name, result))
+            for fn_name in MicroSeries.VECTOR_FUNCTIONS:
+
+                def get_vector_fn(name, res):
+                    def fn(*args, **kwargs):
+                        results = {}
+                        for col in res.numeric_columns:
+                            try:
+                                results[col] = getattr(
+                                    getattr(res, col), name
+                                )(*args, **kwargs)
+                            except Exception:
+                                pass
+                        # Return plain DataFrame - aggregated results don't
+                        # have per-row weights (weights were already applied)
+                        return (
+                            pd.DataFrame(results)
+                            if results
+                            else pd.DataFrame()
+                        )
+
+                    return fn
+
+                setattr(result, fn_name, get_vector_fn(fn_name, result))
+            return result

microdf/tests/test_microseries_dataframe.py

@@ -287,3 +287,60 @@ def test_reset_index_inplace() -> None:
     assert "second" in mdf_multi.columns
     assert list(mdf_multi.index) == [0, 1, 2, 3]
     np.testing.assert_array_equal(mdf_multi.weights.values, weights)
+
+
+def test_loc_preserves_weights() -> None:
+    """Test that .loc[] returns MicroDataFrame with proper weights (issue
+    #265)."""
+    df = mdf.MicroDataFrame(
+        {"one": [1, 1, 1, 1, 1]}, weights=[10, 20, 30, 40, 50]
+    )
+
+    # Filter all rows (should get same weights)
+    filtered = df.loc[df.one == 1]
+    assert isinstance(filtered, MicroDataFrame)
+    assert filtered.one.sum() == 150.0  # Weighted sum
+
+    # Partial filter
+    df2 = mdf.MicroDataFrame(
+        {"x": [1, 2, 3, 4, 5]}, weights=[10, 20, 30, 40, 50]
+    )
+    subset = df2.loc[df2.x > 2]
+    assert isinstance(subset, MicroDataFrame)
+    assert subset.x.sum() == 500.0  # 3*30 + 4*40 + 5*50 = 500
+    np.testing.assert_array_equal(subset.weights.values, [30.0, 40.0, 50.0])
+
+
+def test_iloc_preserves_weights() -> None:
+    """Test that .iloc[] returns MicroDataFrame with proper weights."""
+    df = mdf.MicroDataFrame(
+        {"x": [1, 2, 3, 4, 5]}, weights=[10, 20, 30, 40, 50]
+    )
+
+    # Select rows by position
+    subset = df.iloc[2:5]
+    assert isinstance(subset, MicroDataFrame)
+    assert subset.x.sum() == 500.0  # 3*30 + 4*40 + 5*50 = 500
+    np.testing.assert_array_equal(subset.weights.values, [30.0, 40.0, 50.0])
+
+
+def test_groupby_column_selection() -> None:
+    """Test that groupby column selection preserves weights (issue #193)."""
+    d = mdf.MicroDataFrame(
+        dict(g=["a", "a", "b"], y=[1, 2, 3]), weights=[4, 5, 6]
+    )
+
+    # Test single column string selection
+    result_str = d.groupby("g")["y"].sum()
+    assert result_str["a"] == 14.0  # 1*4 + 2*5 = 14
+    assert result_str["b"] == 18.0  # 3*6 = 18
+
+    # Test list column selection
+    result_list = d.groupby("g")[["y"]].sum()
+    assert result_list.loc["a", "y"] == 14.0
+    assert result_list.loc["b", "y"] == 18.0
+
+    # Aggregated results should be plain DataFrame (no spurious weight column)
+    result_all = d.groupby("g").sum()
+    assert "weight" not in result_all.columns
+    assert list(result_all.columns) == ["y"]

{microdf_python-1.0.2.dist-info → microdf_python-1.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: microdf-python
-Version: 1.0.2
+Version: 1.1.0
 Summary: Weighted pandas DataFrames and Series for survey microdata
 Author-email: Max Ghenis <max@policyengine.org>
 License: MIT
@@ -20,8 +20,6 @@ Requires-Dist: linecheck; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pytest-cov; extra == "dev"
 Requires-Dist: setuptools; extra == "dev"
-Provides-Extra: docs
-Requires-Dist: jupyter_book; extra == "docs"
 Dynamic: license-file
 
 [![Build](https://github.com/PolicyEngine/microdf/workflows/Pull%20request/badge.svg)](https://github.com/PolicyEngine/microdf/actions)

microdf_python-1.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+microdf/__init__.py,sha256=ldRvhE8t4cD7qlxQ5yLndMnWEUYpDw4m-jCDV_psXkY,319
+microdf/microdataframe.py,sha256=X91fUvDS9xMzj2FFAoa6g9EhR_mJv60DXtvK4iIlMrM,33237
+microdf/microseries.py,sha256=MFBStp1IaNVABaf3_Ap_VwXSC3bA8V7rXD4TEG_091g,22452
+microdf/tests/conftest.py,sha256=u-EMyX1-u_nM-YO0RJYCzYHQDXxUI2WQE6GkyJlErqg,150
+microdf/tests/test_microseries_dataframe.py,sha256=U33uBNW6uiF8c7ygDJjwYX_tJJc_i7OKHi-O746XQLY,11398
+microdf_python-1.1.0.dist-info/licenses/LICENSE,sha256=uPs-ASYnzlldpf2z8jeRgQFeEH3FLhSuX0rw0OKWoDU,1067
+microdf_python-1.1.0.dist-info/METADATA,sha256=vPKXByn0atLs1zxTwc6XqESm7iE1wQ4lNtVIIAPBfQA,2420
+microdf_python-1.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+microdf_python-1.1.0.dist-info/top_level.txt,sha256=T2WFPTygQQMdS3GF8YpZ12DKfMGrspbZ3r7z-e3KfiM,8
+microdf_python-1.1.0.dist-info/RECORD,,

microdf_python-1.0.2.dist-info/RECORD

@@ -1,10 +0,0 @@
-microdf/__init__.py,sha256=ldRvhE8t4cD7qlxQ5yLndMnWEUYpDw4m-jCDV_psXkY,319
-microdf/microdataframe.py,sha256=TU4MPHlGSICibPgGkATJYDpVBQ2NNXnyBAqw3kjteeI,25771
-microdf/microseries.py,sha256=MFBStp1IaNVABaf3_Ap_VwXSC3bA8V7rXD4TEG_091g,22452
-microdf/tests/conftest.py,sha256=u-EMyX1-u_nM-YO0RJYCzYHQDXxUI2WQE6GkyJlErqg,150
-microdf/tests/test_microseries_dataframe.py,sha256=2pm_fLZPbv-ccDKnEyqIiw-_pTZwt9J2lnM5YsZHzx4,9401
-microdf_python-1.0.2.dist-info/licenses/LICENSE,sha256=uPs-ASYnzlldpf2z8jeRgQFeEH3FLhSuX0rw0OKWoDU,1067
-microdf_python-1.0.2.dist-info/METADATA,sha256=sOLX9w2xCHwyNeLcW7BpXu6iiTTeaTmxNdMsdJpmD7s,2486
-microdf_python-1.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-microdf_python-1.0.2.dist-info/top_level.txt,sha256=T2WFPTygQQMdS3GF8YpZ12DKfMGrspbZ3r7z-e3KfiM,8
-microdf_python-1.0.2.dist-info/RECORD,,