Repository: PSLmodels/microdf

Issue:
MicroDataFrame[[cols]] produces KeyError
This works in `pandas` and should produce a `MicroDataFrame` with only the column `x`:
```
d = mdf.MicroDataFrame(dict(x=[1, 2], y=[3, 4]), weights=[5, 6])
d[["x"]]
```
But it errors out instead:
```
KeyError: "None of [Index(['x'], dtype='object')] are in the [index]"
```

Discussion:

Patch:
diff --git a/microdf/generic.py b/microdf/generic.py
index f6a91ad..32d688c 100644
--- a/microdf/generic.py
+++ b/microdf/generic.py
@@ -1,8 +1,8 @@
-import numpy as np
-import pandas as pd
 from typing import Callable, Union
 from functools import wraps
 import warnings
+import numpy as np
+import pandas as pd
 
 
 class MicroSeries(pd.Series):
@@ -259,6 +259,16 @@ class MicroSeries(pd.Series):
         gb.weights = pd.Series(self.weights).groupby(*args, **kwargs)
         return gb
 
+    def copy(self, deep=True):
+        res = super().copy(deep)
+        res = MicroSeries(res, weights=self.weights.copy(deep))
+        return res
+
+    def equals(self, other) -> bool:
+        equal_values = super().equals(other)
+        equal_weights = self.weights.equals(other.weights)
+        return equal_values and equal_weights
+
     def __getitem__(self, key):
         result = super().__getitem__(key)
         if isinstance(result, pd.Series):
@@ -593,7 +603,7 @@ class MicroDataFrame(pd.DataFrame):
     def __getitem__(self, key):
         result = super().__getitem__(key)
         if isinstance(result, pd.DataFrame):
-            weights = self.weights.__getitem__(key)
+            weights = self.weights
             return MicroDataFrame(result, weights=weights)
         return result
 
@@ -611,6 +621,19 @@ class MicroDataFrame(pd.DataFrame):
         res = MicroDataFrame(res, weights=self.weights)
         return res
 
+    def copy(self, deep=True):
+        res = super().copy(deep)
+        # This changes the original columns to Series. Undo it:
+        for col in self.columns:
+            self[col] = MicroSeries(self[col])
+        res = MicroDataFrame(res, weights=self.weights.copy(deep))
+        return res
+
+    def equals(self, other) -> bool:
+        equal_values = super().equals(other)
+        equal_weights = self.weights.equals(other.weights)
+        return equal_values and equal_weights
+
     def groupby(self, by: Union[str, list], *args, **kwargs):
         """Returns a GroupBy object with MicroSeriesGroupBy objects for each column
 


Test patch:
diff --git a/microdf/tests/test_generic.py b/microdf/tests/test_generic.py
index aadd6b0..26f57ab 100644
--- a/microdf/tests/test_generic.py
+++ b/microdf/tests/test_generic.py
@@ -152,3 +152,28 @@ def test_rank():
 
     s = mdf.MicroSeries([2, 1, 3], weights=[5, 4, 6])
     assert np.array_equal(s.rank().values, [9, 4, 15])
+
+
+def test_copy_equals():
+    d = mdf.MicroDataFrame(
+        {"x": [1, 2], "y": [3, 4], "z": [5, 6]}, weights=[7, 8]
+    )
+    d_copy = d.copy()
+    d_copy_diff_weights = d_copy.copy()
+    d_copy_diff_weights.weights *= 2
+    assert d.equals(d_copy)
+    assert not d.equals(d_copy_diff_weights)
+    # Same for a MicroSeries.
+    assert d.x.equals(d_copy.x)
+    assert not d.x.equals(d_copy_diff_weights.x)
+
+
+def test_subset():
+    df = mdf.MicroDataFrame(
+        {"x": [1, 2], "y": [3, 4], "z": [5, 6]}, weights=[7, 8]
+    )
+    df_no_z = mdf.MicroDataFrame({"x": [1, 2], "y": [3, 4]}, weights=[7, 8])
+    assert df[["x", "y"]].equals(df_no_z)
+    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)