diff --git a/eland/dataframe.py b/eland/dataframe.py index 54c3fc08..b92f728c 100644 --- a/eland/dataframe.py +++ b/eland/dataframe.py @@ -25,7 +25,7 @@ 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 @@ def drop( 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 diff --git a/eland/tests/dataframe/test_dir_pytest.py b/eland/tests/dataframe/test_dir_pytest.py new file mode 100644 index 00000000..1ab1dc70 --- /dev/null +++ b/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 diff --git a/eland/tests/test_utils_pytest.py b/eland/tests/test_utils_pytest.py new file mode 100644 index 00000000..b9048357 --- /dev/null +++ b/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+") diff --git a/eland/utils.py b/eland/utils.py index 92025720..7b49ae8e 100644 --- a/eland/utils.py +++ b/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 wrapped(*args, **kwargs): 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_]*$" + )