diff --git a/microdf/generic.py b/microdf/generic.py index 4b31aba..b0906cf 100644 --- a/microdf/generic.py +++ b/microdf/generic.py @@ -373,6 +373,11 @@ def __neg__(self, other): def __pos__(self, other): return MicroSeries(super().__pos__(other), weights=self.weights) + def __repr__(self): + return pd.DataFrame( + dict(value=self.values, weight=self.weights.values) + ).__repr__() + MicroSeries.SCALAR_FUNCTIONS = [ fn @@ -620,7 +625,10 @@ def set_weight_col(self, column: str) -> None: def __getitem__(self, key): result = super().__getitem__(key) if isinstance(result, pd.DataFrame): - weights = self.weights + try: + weights = self.weights[key] + except Exception: + weights = self.weights return MicroDataFrame(result, weights=weights) return result @@ -755,3 +763,8 @@ def poverty_count( """ in_poverty = income < threshold return in_poverty.sum() + + def __repr__(self): + df = pd.DataFrame(self) + df["weight"] = self.weights + return df[[df.columns[-1]] + list(df.columns[:-1])].__repr__() diff --git a/microdf/tests/test_generic.py b/microdf/tests/test_generic.py index 3dfaffa..3a1a58e 100644 --- a/microdf/tests/test_generic.py +++ b/microdf/tests/test_generic.py @@ -200,3 +200,9 @@ def test_subset(): df_no_z_diff_weights = df_no_z.copy() df_no_z_diff_weights.weights += 1 assert not df[["x", "y"]].equals(df_no_z_diff_weights) + + +def test_value_subset(): + d = mdf.MicroDataFrame({"x": [1, 2, 3], "y": [1, 2, 2]}, weights=[4, 5, 6]) + d2 = d[d.y > 1] + assert d2.y.shape == d2.weights.shape