o 8Di;@sddlmZddlZddlZddlZddlZddlZddlZddl m Z m Z ddl m Z ddlmZddlmZddlmZddlmZejrcdd lmZdd lmZmZmZdd lmZdd lmZe d Z!ej"ddZ#e$dZ%GdddeZ&Gddde'Z(dddZ)dd"d#Z*ej+dd'd(Z,ej+dd*d(Z,ej+dd,d(Z,d-d(Z,ej+dd.d/Z-ej+dd0d/Z-d1d/Z-d2d3dd6d7Z. 8ddd>d?Z/ dddGdHZ0ddJdKZ1ddQdRZ2ddUdVZ3ddYdZZ4e dd^d_Z5ddcddZ6ddgdhZ7ddidjZ8ddmdnZ9ddqdrZ:ddudvZ; wddd{d|ZdddZ?dddZ@dddZAdddZBdddZCdddZDhdZEdddZFeGdZHeGdZIGdddejJeHeIfZKdS)) annotationsN) CollectionSet)contextmanagercopy)get_close_matches)Enumcountexp)AET) DialectType Expressionz(?rrrr7?s   r7valuet.Collection[T] t.List[T]cCdSrrrBrrr ensure_listGrrGt.ListcCrErrrFrrrrGKrrcCrErrrFrrrrGOrcCs(|durgSt|ttfrt|S|gS)z Ensures that a value is a list, otherwise casts or wraps it into one. Args: value: The value of interest. Returns: The value cast as a list if it's a list or a tuple, or else the value wrapped in a list. N) isinstancelisttuplerFrrrrGSs  cCrErrrFrrrensure_collectionerrLcCrErrrFrrrrLircCs.|durgSt|trt|ttfs|S|gS)z Ensures that a value is a collection (excluding `str` and `bytes`), otherwise wraps it into a list. Args: value: The value of interest. Returns: The value if it's a collection, or else the value wrapped in a list. N)rIrr.bytesrFrrrrLms  z, )separgsrNcGs|dd|DS)z Formats any number of string arguments as CSV. Args: args: The string arguments to format. sep: The argument separator. Returns: The arguments formatted as a CSV string. css|]}|r|VqdSrr.0argrrr szcsv..)join)rNrOrrrcsv~s rUr module_nameclassest.Type | t.Tuple[t.Type, ...]excludet.List[t.Type]cs&ddttj|fddDS)ac Returns all subclasses for a collection of classes, possibly excluding some of them. Args: module_name: The name of the module to search for subclasses in. classes: Class(es) we want to find the subclasses of. exclude: Class(es) we want to exclude from the returned list. Returns: The target subclasses. cSsg|]\}}|qSrr)rQ_r%rrr szsubclasses..cst|o t|o |vSr)inspectisclass issubclass)r%rWrYrrszsubclasses..)r] getmemberssysmodules)rVrWrYrr`r subclassess  rethisexp.Expression expressions t.List[E]offsetdialectrcCs|rt|dkr |S|d}ddlm}ddlm}ddlm}|js)|||dt |j |jj |j j j |j j jfvr>|S|jsG|||dt |j |jj |j jvrctd||||}|gS|S)a Applies an offset to a given integer literal expression. Args: this: The target of the index. expressions: The expression the offset will be applied to, wrapped in a list. offset: The offset that will be applied. dialect: the dialect of interest. Returns: The original expression with the offset applied to it, wrapped in a list. If the provided `expressions` argument contains more than one expression, it's returned unaffected. r3rr )annotate_types)simplify)rkz Applying array index offset (%s))lenrr sqlglot.optimizer.annotate_typesrlsqlglot.optimizer.simplifyrmtypetcastDataTyperfTypeUNKNOWNARRAY INTEGER_TYPESloggerinfo)rfrhrjrk expressionr rlrmrrrapply_index_offsets(       r|rcCstd|S)zDConverts `name` from camelCase to snake_case and returns the result.r[)CAMEL_CASE_PATTERNsubupperrrrrcamel_to_snake_casesrr{rfunct.Callable[[Expression], E]rcCsd} |durtt|D]}t||_qt|}||}t|}t|D] }d|_t||_q)t|}||krG|D]}d|_q> |Sq)z Applies a transformation to a given expression until a fix point is reached. Args: expression: The expression to be transformed. func: The transformation to be applied. Returns: The transformed expression. N)reversedrKwalkhash_hash)r{rend_hashr4 start_hashexpression_nodesrrrwhile_changings$     rdagt.Dict[T, t.Set[T]]cCsg}t|D]\}}|D] }||vrt||<qq|rJdd|D}|s,td|D]}||q.|D]}||8}q:|t||s|S)z Sorts a given directed acyclic graph in topological order. Args: dag: The graph to be sorted. Returns: A list that contains all of the graph's nodes in topological order. cSsh|]\}}|s|qSrr)rQnodedepsrrr sztsort..z Cycle error)rKitemssetr8popvaluesextendsorted)rresultrrdepcurrentrrrtsorts$     r file_namet.TextIOcCsbt|d}|ddk}Wdn1swY|r*ddl}|j|dddSt|d dd S) zSOpen a file that may be compressed as gzip and return it in universal newline mode.rbrsNrrtr5)newlinezutf-8)encodingr)openreadgzip)rfgzippedrrrr open_file%s rread_csv exp.ReadCSVr&ccsz|j}t|j}d}tdd|D}t||D] \}}|dkr#|}qzddl}|j||dVW|dS|w)z Returns a csv reader given the expression `READ_CSV(name, ['delimiter', '|', ...])`. Args: read_csv: A `ReadCSV` function call. Yields: A python csv reader. ,css|]}|jVqdSrrrPrrrrSAszcsv_reader.. delimiterrN)r)rhrriterziprUreaderclose)rrOfilerkvcsv_rrr csv_reader2s rtakent.Collection[str]basecCsH||vr|Sd}|d|}||vr"|d7}|d|}||vs|S)z Searches for a new name. Args: taken: A collection of taken names. base: Base name to alter. Returns: The new, available name. rr[r3r)rrinewrrr find_new_nameNs rtextboolcC t|tSr)is_typer?rrrris_inte rcCrr)rfloatrrrris_floatirr target_typet.TypecCs$z||WdStyYdSwNTF)r8)rrrrrrms  rprefixt.Callable[[], str]cstfddS)zTReturns a name generator given a prefix (e.g. a0, a1, a2, ... if the prefix is "a").cstSrnextrrsequencerrraxszname_sequence..r )rrrr name_sequenceusrr%t.DictcKsiddt|D|S)z9Returns a dictionary created from an object's attributes.cSs,i|]\}}|t|dr|nt|qSr)hasattrr)rQrrrrr ~s,z"object_to_dict..)varsr)r%kwargsrrrobject_to_dict{s rT min_num_wordsfill_from_startt.List[t.Optional[str]]cCs:||}|rdg|t||S|dg|t|S)a 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'] Returns: The list of words returned by `split`, possibly augmented by a number of `None` values. N)splitrn)rBrNrrwordsrrrsplit_num_wordss rcCs(ddlm}t|dot|tt|f S)aF Checks if the value is an iterable, excluding the types `str` and `bytes`. Examples: >>> is_iterable([1,2]) True >>> is_iterable("test") False Args: value: The value to check if it is an iterable. Returns: A `bool` value indicating if it is an iterable. rr__iter__)rrrrIr.rM)rBrrrr is_iterables rr%t.Iterable[t.Iterable[t.Any] | t.Any]t.Iterator[t.Any]ccs.|D]}t|rt|EdHq|VqdS)a Flattens an iterable that can contain both iterable and non-iterable elements. Objects of type `str` and `bytes` are not regarded as iterables. Examples: >>> list(flatten([[1, 2], 3, {4}, (5, "bla")])) [1, 2, 3, 4, 5, 'bla'] >>> list(flatten([1, 2, 3])) [1, 2, 3] Args: values: The value to be flattened. Yields: Non-iterable elements in `values`. N)rflatten)rrBrrrrs rdcCsBz dttt|WStyYdSty YdSw)a Get the nesting depth of a dictionary. Example: >>> dict_depth(None) 0 >>> dict_depth({}) 1 >>> dict_depth({"a": "b"}) 1 >>> dict_depth({"a": {}}) 2 >>> dict_depth({"a": {"b": {}}}) 3 r3r) dict_depthrrrAttributeError StopIteration)rrrrrs  rit t.Iterable[T]cCstdd|DS)z=Returns the first element from an iterable (useful for sets).css|]}|VqdSrr)rQrrrrrSszfirst..r)rrrrfirstsrt.Optional[str | bool]cCs:t|ts |dur |S|}|dvrdS|dvrdS|S)N)true1T)false0F)rIrlower)rB value_lowerrrrto_boolsrrangest.List[t.Tuple[A, A]]cCsj|sgSt|}|dg}|ddD]\}}|d\}}||kr+|t||f|d<q|||fq|S)z Merges a sequence of ranges, represented as tuples (low, high) whose values belong to some totally-ordered set. Example: >>> merge_ranges([(1, 3), (2, 6)]) [(1, 6)] rr3N)rmaxappend)rmergedstartend last_startlast_endrrr merge_rangess   rcCs(z tj|WdStyYdSwr)datetimedate fromisoformatr8rrrr is_iso_date   rcCs(z tj|WdStyYdSwr)rrr8rrrris_iso_datetimerr>dayweekyearmonthquarter year_montht.Optional[exp.Expression]cCs|duo |jtvSr)rr DATE_UNITS)r{rrr is_date_unit)srKVc@s8eZdZdZdddZdd d ZdddZdddZdS)SingleValuedMappingz Mapping where all keys return the same value. This rigamarole is meant to avoid copying keys, which was originally intended as an optimization while qualifying columns for tables with lots of columns. keyst.Collection[K]rBrcCs"t|tr|nt||_||_dSr)rIrr_keys_value)r,r rBrrr__init__9s zSingleValuedMapping.__init__keyrr(cCs||jvr|jSt|r)r r KeyError)r,rrrr __getitem__=s zSingleValuedMapping.__getitem__r?cC t|jSr)rnr r,rrr__len__BrzSingleValuedMapping.__len__ t.Iterator[K]cCrr)rr rrrrrErzSingleValuedMapping.__iter__N)r r rBr)rrr(r)r(r?)r(r)r r!r"r#rrrrrrrrr 1s    r )r-r.r/r.r0r1r(r2)r<r=r>r?r(r@)rBrCr(rD)rBr2r(rH)rBrr(rD)rBrCr(rC)rBrr(rC)rOr.rNr.r(r.)r)rVr.rWrXrYrXr(rZr) rfrgrhrirjr?rkrr(ri)rr.r(r.)r{rrrr(r)rrr(rD)rr.r(r)rrr(r&)rrrr.r(r.)rr.r(r)rr.rrr(r)rr.r(r)r%r&r(r)T) rBr.rNr.rr?rrr(r)rBr&r(r)rrr(r)rrr(r?)rrr(r)rBrr(r)rrr(r)r{rr(r)L __future__rrr]loggingrerctypingrrcollections.abcrr contextlibrrdifflibrenumr itertoolsr TYPE_CHECKINGrr sqlglot._typingrrrsqlglot.dialects.dialectrsqlglot.expressionsrcompiler} version_infoPYTHON_VERSION getLoggerryrpropertyr$r;r7overloadrGrLrUrer|rrrrrrrrrrrrrrrrrrrrrrTypeVarrrMappingr rrrrs                   0  ) "