Merge pull request #219 from nikhilwoodruff/bug-fixes

Bug fixes

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__()

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