Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dataframe column names autocompletion #223

Merged
merged 1 commit into from Jul 2, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
12 changes: 11 additions & 1 deletion eland/dataframe.py
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down
20 changes: 20 additions & 0 deletions 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
22 changes: 22 additions & 0 deletions 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+")
10 changes: 10 additions & 0 deletions eland/utils.py
Expand Up @@ -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
Expand All @@ -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_]*$"
)