diff options
Diffstat (limited to 'sqlglot/helper.py')
-rw-r--r-- | sqlglot/helper.py | 53 |
1 files changed, 53 insertions, 0 deletions
diff --git a/sqlglot/helper.py b/sqlglot/helper.py index c4dd91e..c3a23d3 100644 --- a/sqlglot/helper.py +++ b/sqlglot/helper.py @@ -2,7 +2,9 @@ import inspect import logging import re import sys +import typing as t from contextlib import contextmanager +from copy import copy from enum import Enum CAMEL_CASE_PATTERN = re.compile("(?<!^)(?=[A-Z])") @@ -162,3 +164,54 @@ def find_new_name(taken, base): i += 1 new = f"{base}_{i}" return new + + +def object_to_dict(obj, **kwargs): + return {**{k: copy(v) for k, v in vars(obj).copy().items()}, **kwargs} + + +def split_num_words(value: str, sep: str, min_num_words: int, fill_from_start: bool = True) -> t.List[t.Optional[str]]: + """ + Perform a split on a value and return N words as a result with None used for words that don't exist. + + Args: + value: The value to be split + sep: The value to use to split on + min_num_words: The minimum number of words that are going to be in the result + fill_from_start: Indicates that if None values should be inserted at the start or end of the list + + Examples: + >>> split_num_words("db.table", ".", 3) + [None, 'db', 'table'] + >>> split_num_words("db.table", ".", 3, fill_from_start=False) + ['db', 'table', None] + >>> split_num_words("db.table", ".", 1) + ['db', 'table'] + """ + words = value.split(sep) + if fill_from_start: + return [None] * (min_num_words - len(words)) + words + return words + [None] * (min_num_words - len(words)) + + +def flatten(values: t.Iterable[t.Union[t.Iterable[t.Any], t.Any]]) -> t.Generator[t.Any, None, None]: + """ + Flattens a list that can contain both iterables and non-iterable elements + + Examples: + >>> list(flatten([[1, 2], 3])) + [1, 2, 3] + >>> list(flatten([1, 2, 3])) + [1, 2, 3] + + Args: + values: The value to be flattened + + Returns: + Yields non-iterable elements (not including str or byte as iterable) + """ + for value in values: + if hasattr(value, "__iter__") and not isinstance(value, (str, bytes)): + yield from flatten(value) + else: + yield value |