Skip to content

Commit

Permalink
#222- Dataframe column names autocompletion
Browse files Browse the repository at this point in the history
  • Loading branch information
leonardbinet committed Jul 1, 2020
1 parent f639410 commit c1dd24a
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 1 deletion.
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(NDFrame, self).__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 not is_valid_attr_name('_')
assert not 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+')
8 changes: 8 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,10 @@ 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_]*$") and s.replace('_', '')

0 comments on commit c1dd24a

Please sign in to comment.