diff options
Diffstat (limited to 'sqlglot/helper.py')
-rw-r--r-- | sqlglot/helper.py | 45 |
1 files changed, 38 insertions, 7 deletions
diff --git a/sqlglot/helper.py b/sqlglot/helper.py index 379c2e7..8c5808d 100644 --- a/sqlglot/helper.py +++ b/sqlglot/helper.py @@ -11,7 +11,8 @@ from copy import copy from enum import Enum if t.TYPE_CHECKING: - from sqlglot.expressions import Expression, Table + from sqlglot import exp + from sqlglot.expressions import Expression T = t.TypeVar("T") E = t.TypeVar("E", bound=Expression) @@ -150,7 +151,7 @@ def apply_index_offset(expressions: t.List[E], offset: int) -> t.List[E]: if expression.is_int: expression = expression.copy() logger.warning("Applying array index offset (%s)", offset) - expression.args["this"] = str(int(expression.args["this"]) + offset) + expression.args["this"] = str(int(expression.this) + offset) return [expression] return expressions @@ -228,19 +229,18 @@ def open_file(file_name: str) -> t.TextIO: @contextmanager -def csv_reader(table: Table) -> t.Any: +def csv_reader(read_csv: exp.ReadCSV) -> t.Any: """ Returns a csv reader given the expression `READ_CSV(name, ['delimiter', '|', ...])`. Args: - table: a `Table` expression with an anonymous function `READ_CSV` in it. + read_csv: a `ReadCSV` function call Yields: A python csv reader. """ - file, *args = table.this.expressions - file = file.name - file = open_file(file) + args = read_csv.expressions + file = open_file(read_csv.name) delimiter = "," args = iter(arg.name for arg in args) @@ -354,3 +354,34 @@ def flatten(values: t.Iterable[t.Iterable[t.Any] | t.Any]) -> t.Generator[t.Any, yield from flatten(value) else: yield value + + +def dict_depth(d: t.Dict) -> int: + """ + Get the nesting depth of a dictionary. + + For example: + >>> dict_depth(None) + 0 + >>> dict_depth({}) + 1 + >>> dict_depth({"a": "b"}) + 1 + >>> dict_depth({"a": {}}) + 2 + >>> dict_depth({"a": {"b": {}}}) + 3 + + Args: + d (dict): dictionary + Returns: + int: depth + """ + try: + return 1 + dict_depth(next(iter(d.values()))) + except AttributeError: + # d doesn't have attribute "values" + return 0 + except StopIteration: + # d.values() returns an empty sequence + return 1 |