add/delete str/dt/cat dynamically from __dir__ (fix for pandas-dev#9627)

Author: mortada

pandas/core/base.py

@@ -86,16 +86,22 @@ def __unicode__(self):
         # Should be overwritten by base classes
         return object.__repr__(self)
 
-    def _local_dir(self):
-        """ provide addtional __dir__ for this object """
-        return []
+    def _dir_additions(self):
+        """ add addtional __dir__ for this object """
+        return set()
+
+    def _dir_deletions(self):
+        """ delete unwanted __dir__ for this object """
+        return set()
 
     def __dir__(self):
         """
         Provide method name lookup and completion
         Only provide 'public' methods
         """
-        return list(sorted(list(set(dir(type(self)) + self._local_dir()))))
+        rv = set(dir(type(self)))
+        rv = (rv - self._dir_deletions()) | self._dir_additions()
+        return sorted(rv)
 
     def _reset_cache(self, key=None):
         """
@@ -498,8 +504,7 @@ def searchsorted(self, key, side='left'):
         #### needs tests/doc-string
         return self.values.searchsorted(key, side=side)
 
-    # string methods
-    def _make_str_accessor(self):
+    def _check_str_accessor(self):
         from pandas.core.series import Series
         from pandas.core.index import Index
         if isinstance(self, Series) and not com.is_object_dtype(self.dtype):
@@ -512,10 +517,24 @@ def _make_str_accessor(self):
         elif isinstance(self, Index) and self.inferred_type != 'string':
             raise AttributeError("Can only use .str accessor with string "
                                  "values (i.e. inferred_type is 'string')")
+
+    # string methods
+    def _make_str_accessor(self):
+        self._check_str_accessor()
         return StringMethods(self)
 
     str = AccessorProperty(StringMethods, _make_str_accessor)
 
+    def _dir_additions(self):
+        return set()
+
+    def _dir_deletions(self):
+        try:
+            self._check_str_accessor()
+        except AttributeError:
+            return set(['str'])
+        return set()
+
     _shared_docs['drop_duplicates'] = (
         """Return %(klass)s with duplicate values removed
 

pandas/core/generic.py

@@ -146,10 +146,10 @@ def __unicode__(self):
         prepr = '[%s]' % ','.join(map(com.pprint_thing, self))
         return '%s(%s)' % (self.__class__.__name__, prepr)
 
-    def _local_dir(self):
+    def _dir_additions(self):
         """ add the string-like attributes from the info_axis """
-        return [c for c in self._info_axis
-                if isinstance(c, string_types) and isidentifier(c)]
+        return set([c for c in self._info_axis
+                if isinstance(c, string_types) and isidentifier(c)])
 
     @property
     def _constructor_sliced(self):

pandas/core/groupby.py

@@ -498,8 +498,8 @@ def _set_result_index_ordered(self, result):
         result.index = self.obj.index
         return result
 
-    def _local_dir(self):
-        return sorted(set(self.obj._local_dir() + list(self._apply_whitelist)))
+    def _dir_additions(self):
+        return self.obj._dir_additions() | self._apply_whitelist
 
     def __getattr__(self, attr):
         if attr in self._internal_names_set:

pandas/core/series.py

@@ -29,6 +29,7 @@
 from pandas.core.internals import SingleBlockManager
 from pandas.core.categorical import Categorical, CategoricalAccessor
 from pandas.tseries.common import (maybe_to_datetimelike,
+                                   is_datetimelike,
                                    CombinedDatetimelikeProperties)
 from pandas.tseries.index import DatetimeIndex
 from pandas.tseries.tdi import TimedeltaIndex
@@ -2517,6 +2518,23 @@ def _make_cat_accessor(self):
 
     cat = base.AccessorProperty(CategoricalAccessor, _make_cat_accessor)
 
+    def _dir_deletions(self):
+        return set(['str', 'cat', 'dt'])
+
+    def _dir_additions(self):
+        rv = set()
+        if is_datetimelike(self):
+            rv.add('dt')
+        elif com.is_categorical_dtype(self.dtype):
+            rv.add('cat')
+        else:
+            try:
+                self._check_str_accessor()
+                rv.add('str')
+            except AttributeError:
+                pass
+        return rv
+
 Series._setup_axes(['index'], info_axis=0, stat_axis=0,
                    aliases={'rows': 0})
 Series._add_numeric_operations()

pandas/tests/test_series.py

@@ -242,6 +242,25 @@ def test_dt_accessor_api(self):
                 s.dt
             self.assertFalse(hasattr(s, 'dt'))
 
+    def test_tab_completion(self):
+        s = Series(list('abcd'))
+        # Series of str values should have .str but not .dt/.cat in __dir__
+        self.assertTrue('str' in s.__dir__())
+        self.assertTrue('dt' not in s.__dir__())
+        self.assertTrue('cat' not in s.__dir__())
+
+        # similiarly for .dt
+        s = Series(date_range('1/1/2015', periods=5))
+        self.assertTrue('dt' in s.__dir__())
+        self.assertTrue('str' not in s.__dir__())
+        self.assertTrue('cat' not in s.__dir__())
+
+        # similiarly for .cat
+        s = Series(list('abbcd'), dtype="category")
+        self.assertTrue('cat' in s.__dir__())
+        self.assertTrue('str' not in s.__dir__())
+        self.assertTrue('dt' not in s.__dir__())
+
     def test_binop_maybe_preserve_name(self):
 
         # names match, preserve