Add column names to DataFrame.__dir__ for better auto-completion support

eland/dataframe.py

@@ -25,7 +25,7 @@ from eland.ndframe import NDFrame
 from eland.series import Series
 from eland.common import DEFAULT_NUM_ROWS_DISPLAYED, docstring_parameter
 from eland.filter import BooleanFilter
-from eland.utils import deprecated_api
+from eland.utils import deprecated_api, is_valid_attr_name
 
 
 class DataFrame(NDFrame):
@@ -445,6 +445,16 @@ class DataFrame(NDFrame):
     def __getitem__(self, key):
         return self._getitem(key)
 
+    def __dir__(self):
+        """
+        Provide autocompletion on field names in interactive environment.
+        """
+        return super().__dir__() + [
+            column_name
+            for column_name in self._query_compiler.columns.to_list()
+            if is_valid_attr_name(column_name)
+        ]
+
     def __repr__(self):
         """
         From pandas

eland/tests/dataframe/test_dir_pytest.py

@@ -0,0 +1,20 @@
+# Licensed to Elasticsearch B.V under one or more agreements.
+# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+# See the LICENSE file in the project root for more information
+
+# File called _pytest for PyCharm compatibility
+
+
+from eland.tests.common import TestData
+
+
+class TestDataFrameDir(TestData):
+    def test_flights_dir(self):
+        ed_flights = self.ed_flights()
+
+        print(dir(ed_flights))
+
+        autocomplete_attrs = dir(ed_flights)
+
+        for c in ed_flights.columns:
+            assert c in autocomplete_attrs

eland/tests/test_utils_pytest.py

@@ -0,0 +1,22 @@
+# Licensed to Elasticsearch B.V under one or more agreements.
+# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+# See the LICENSE file in the project root for more information
+
+# File called _pytest for PyCharm compatibility
+
+from eland.utils import is_valid_attr_name
+
+
+class TestUtils:
+    def test_is_valid_attr_name(self):
+        assert is_valid_attr_name("_piZZa")
+        assert is_valid_attr_name("nice_pizza_with_2_mushrooms")
+        assert is_valid_attr_name("_2_pizze")
+        assert is_valid_attr_name("_")
+        assert is_valid_attr_name("___")
+
+        assert not is_valid_attr_name("4")
+        assert not is_valid_attr_name(4)
+        assert not is_valid_attr_name(None)
+        assert not is_valid_attr_name("4pizze")
+        assert not is_valid_attr_name("pizza+")

eland/utils.py

@@ -2,6 +2,7 @@
 # Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
 # See the LICENSE file in the project root for more information
 
+import re
 import functools
 import warnings
 from typing import Callable, TypeVar
@@ -23,3 +24,12 @@ def deprecated_api(replace_with: str) -> Callable[[F], F]:
         return wrapped
 
     return wrapper
+
+
+def is_valid_attr_name(s):
+    """
+    Ensure the given string can be used as attribute on an object instance.
+    """
+    return isinstance(s, str) and re.search(
+        string=s, pattern=r"^[a-zA-Z_][a-zA-Z0-9_]*$"
+    )