Merge pull request #12 from jseabold/gh-8

Support dt property comparisons and between

Author: jseabold

pandas_selectable/__init__.py

@@ -2,6 +2,7 @@
 import operator
 from functools import wraps
 
+import numpy as np
 import pandas as pd
 from pandas.core.accessor import CachedAccessor
 from pandas.core.indexes.accessors import (
@@ -70,31 +71,6 @@ def _wrap_result(self, *args, **kwargs):
         return self._parent_frame.loc[bool_idx]
 
 
-class SelectDatetimeProperties(DatetimeProperties):
-    def __init__(self, parent, *args, **kwargs):
-        # datetime properties holds an attribute _parent
-        # we need to add the parent_frame (or series) to the subclass instances
-        self._parent_frame = parent
-        super().__init__(*args, **kwargs)
-
-    def __getattribute__(self, attr):
-        if (
-            not attr.startswith("_")
-            and inspect.isroutine(  # noqa
-                getattr(DatetimeProperties, attr, None)
-            )
-            and attr not in _date_boolean_methods
-        ):  # noqa
-            raise NotImplementedError(
-                "Boolean selection with this method " "does not make sense."
-            )
-        elif attr in _date_boolean_methods:
-            idx = super().__getattribute__(attr)
-            return self._parent_frame.loc[idx]
-        else:
-            return super().__getattribute__(attr)
-
-
 class SelectPeriodProperties(PeriodProperties):
     def __init__(self, parent, *args, **kwargs):
         self._parent_frame = parent
@@ -207,6 +183,7 @@ class SelectableColumn:
     notna = selector_wrapper(pd.Series, "notna")
     notnull = selector_wrapper(pd.Series, "notnull")
     isin = selector_wrapper(pd.Series, "isin")
+    between = selector_wrapper(pd.Series, "between")
 
     def __init__(self, parent, series=None):
         # if accessed as the series accessor, parent is the series
@@ -252,3 +229,37 @@ def __getitem__(self, key):
     @property
     def index(self):
         return SelectableIndex(self._frame)
+
+
+class SelectDatetimeProperties(DatetimeProperties):
+    def __init__(self, parent, *args, **kwargs):
+        # datetime properties holds an attribute _parent
+        # we need to add the parent_frame (or series) to the subclass instances
+        self._parent_frame = parent
+        super().__init__(*args, **kwargs)
+
+    def __getattribute__(self, attr):
+        if (
+            not attr.startswith("_")
+            and inspect.isroutine(  # noqa
+                getattr(DatetimeProperties, attr, None)
+            )
+            and attr not in _date_boolean_methods
+        ):  # noqa
+            raise NotImplementedError(
+                "Boolean selection with this method " "does not make sense."
+            )
+        elif attr in _date_boolean_methods:
+            idx = super().__getattribute__(attr)
+            return self._parent_frame.loc[idx]
+        else:
+            got_attr = super().__getattribute__(attr)
+            # this allows things like dt.day, dt.month to be selectable
+            # for the parent frame. assumes they're all properties.
+            if (
+                isinstance(got_attr, pd.Series)
+                and not attr.startswith('_')
+                and isinstance(getattr(self.__class__, attr), property)
+            ):
+                return SelectableColumn(self._parent_frame, got_attr)
+            return got_attr

pandas_selectable/tests/test_select.py

@@ -56,6 +56,9 @@ def test_dt(dta):
     assert dta.select.B.dt.is_quarter_start.equals(dta.loc[[0]])
     assert len(dta.select.B.dt.is_year_end) == 0
     assert dta.select.B.dt.is_year_start.equals(dta.loc[[0]])
+    assert len(dta.select.B.dt.day == 15) == 1
+    assert len(dta.select.B.dt.month == 2) == 0
+    assert len(dta.select.B.between("2020-01-03", "2020-01-05")) == 3
 
 
 def test_index(dta):