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):
         """
@@ -516,6 +522,16 @@ def _make_str_accessor(self):
 
     str = AccessorProperty(StringMethods, _make_str_accessor)
 
+    def _dir_additions(self):
+        return set()
+
+    def _dir_deletions(self):
+        try:
+            getattr(self, 'str')
+        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
@@ -2521,6 +2522,21 @@ def _make_cat_accessor(self):
 
     cat = base.AccessorProperty(CategoricalAccessor, _make_cat_accessor)
 
+    def _dir_deletions(self):
+        return self._accessors
+
+    def _dir_additions(self):
+        rv = set()
+        # these accessors are mutually exclusive, so break loop when one exists
+        for accessor in self._accessors:
+            try:
+                getattr(self, accessor)
+                rv.add(accessor)
+                break
+            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_index.py

@@ -1232,6 +1232,13 @@ def test_str_attribute(self):
         expected = Series(range(2), index=['a1', 'a2'])
         tm.assert_series_equal(s[s.index.str.startswith('a')], expected)
 
+    def test_tab_completion(self):
+        idx = Index(list('abcd'))
+        self.assertTrue('str' in idx.__dir__())
+
+        idx = Index(range(4))
+        self.assertTrue('str' not in idx.__dir__())
+
     def test_indexing_doesnt_change_class(self):
         idx = Index([1, 2, 3, 'a', 'b', 'c'])
 

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