Allow selecting based on dt.day, etc. comparisons

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
@@ -252,3 +228,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