Skip to content

Commit

Permalink
Fix bug #925 by ensuring kwargs are lists and handling edge cases
Browse files Browse the repository at this point in the history
Signed-off-by: Sepehr-A <sepehr@stardustsystems.net>
  • Loading branch information
Sepehr-A committed Apr 4, 2024
1 parent b22700e commit 8ed3ab9
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 1 deletion.
3 changes: 2 additions & 1 deletion suzieq/engines/pandas/topology.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ def __init__(self, baseobj):
self.lsdb = pd.DataFrame()
self._a_df = pd.DataFrame()
self._ip_table = pd.DataFrame()
self._namespaces = []

@staticmethod
def table_name():
Expand Down Expand Up @@ -77,7 +78,7 @@ class Services:
extra_cols: list
augment: any

self._namespaces = kwargs.get("namespace", self.ctxt.namespace)
self._namespaces = kwargs.get("namespace", [])
hostname = kwargs.pop('hostname', [])
user_query = kwargs.pop('query_str', '')
polled = kwargs.pop('polled', '')
Expand Down
23 changes: 23 additions & 0 deletions suzieq/sqobjects/basicobj.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,8 @@ def get(self, **kwargs) -> pd.DataFrame:
elif isinstance(val, str):
kwargs[k] = v(val)

# This raises TypeError if it fails
self._validate_list_args(**kwargs)
result = self.engine.get(**kwargs, columns=columns)
if self._is_result_empty(result):
fields = self._get_empty_cols(columns, 'get', **kwargs)
Expand All @@ -245,6 +247,8 @@ def summarize(self, **kwargs) -> pd.DataFrame:
raise AttributeError('No analysis engine specified')

self.validate_summarize_input(**kwargs)
# This raises TypeError if it fails
self._validate_list_args(**kwargs)

result = self.engine.summarize(**kwargs)
if self._is_result_empty(result):
Expand Down Expand Up @@ -284,6 +288,8 @@ def aver(self, **kwargs):
'''Assert one or more checks on table'''
kwargs.pop('ignore_warning', None)
columns = kwargs.pop('columns', self.columns)
# This raises TypeError if it fails
self._validate_list_args(**kwargs)
if self._valid_assert_args:
result = self._assert_if_supported(**kwargs, columns=columns)
if self._is_result_empty(result):
Expand Down Expand Up @@ -312,6 +318,8 @@ def top(self, what: str = '', count: int = 5, reverse: bool = False,
df = pd.DataFrame({'error': [f'{error}']})
return df

# This raises TypeError if it fails
self._validate_list_args(**kwargs)
# This raises ValueError if it fails
table_schema = SchemaForTable(self._table, self.all_schemas)
if not self._field_exists(table_schema, what):
Expand Down Expand Up @@ -573,3 +581,18 @@ def _empty_result(self, columns: List[str]) -> pd.DataFrame:
pd.DataFrame: empty dataframe
"""
return pd.DataFrame(columns=columns)

def _validate_list_args(self, **kwargs):
"""Method to ensure specified kwargs are lists."""
# Specify which arguments should be lists
list_args = ['namespace', 'hostname']
for arg in list_args:
if arg not in kwargs or kwargs[arg] is None:
continue
if not isinstance(kwargs[arg], list):
error_msg = (
f"The argument '{arg}' must be a List of strings, "
f"got '{type(kwargs[arg]).__name__}' instead. current "
f"val: {kwargs[arg]}"
)
raise TypeError(error_msg)

0 comments on commit 8ed3ab9

Please sign in to comment.