Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes15kdownloads
sql.py664 linesDownload Raw Back to sqlparse
1#2# Copyright (C) 2009-2020 the sqlparse authors and contributors3# <see AUTHORS file>4#5# This module is part of python-sqlparse and is released under6# the BSD License: https://opensource.org/licenses/BSD-3-Clause7 8"""This module contains classes representing syntactical elements of SQL."""9 10import re11 12from sqlparse import tokens as T13from sqlparse.exceptions import SQLParseError14from sqlparse.utils import imt, remove_quotes15 16 17class NameAliasMixin:18    """Implements get_real_name and get_alias."""19 20    def get_real_name(self):21        """Returns the real name (object name) of this identifier."""22        # a.b23        dot_idx, _ = self.token_next_by(m=(T.Punctuation, '.'))24        return self._get_first_name(dot_idx, real_name=True)25 26    def get_alias(self):27        """Returns the alias for this identifier or ``None``."""28 29        # "name AS alias"30        kw_idx, kw = self.token_next_by(m=(T.Keyword, 'AS'))31        if kw is not None:32            return self._get_first_name(kw_idx + 1, keywords=True)33 34        # "name alias" or "complicated column expression alias"35        _, ws = self.token_next_by(t=T.Whitespace)36        if len(self.tokens) > 2 and ws is not None:37            return self._get_first_name(reverse=True)38 39 40class Token:41    """Base class for all other classes in this module.42 43    It represents a single token and has two instance attributes:44    ``value`` is the unchanged value of the token and ``ttype`` is45    the type of the token.46    """47 48    __slots__ = ('value', 'ttype', 'parent', 'normalized', 'is_keyword',49                 'is_group', 'is_whitespace')50 51    def __init__(self, ttype, value):52        value = str(value)53        self.value = value54        self.ttype = ttype55        self.parent = None56        self.is_group = False57        self.is_keyword = ttype in T.Keyword58        self.is_whitespace = self.ttype in T.Whitespace59        self.normalized = value.upper() if self.is_keyword else value60 61    def __str__(self):62        return self.value63 64    # Pending tokenlist __len__ bug fix65    # def __len__(self):66    #     return len(self.value)67 68    def __repr__(self):69        cls = self._get_repr_name()70        value = self._get_repr_value()71 72        q = '"' if value.startswith("'") and value.endswith("'") else "'"73        return "<{cls} {q}{value}{q} at 0x{id:2X}>".format(74            id=id(self), **locals())75 76    def _get_repr_name(self):77        return str(self.ttype).split('.')[-1]78 79    def _get_repr_value(self):80        raw = str(self)81        if len(raw) > 7:82            raw = raw[:6] + '...'83        return re.sub(r'\s+', ' ', raw)84 85    def flatten(self):86        """Resolve subgroups."""87        yield self88 89    def match(self, ttype, values, regex=False):90        """Checks whether the token matches the given arguments.91 92        *ttype* is a token type. If this token doesn't match the given token93        type.94        *values* is a list of possible values for this token. The values95        are OR'ed together so if only one of the values matches ``True``96        is returned. Except for keyword tokens the comparison is97        case-sensitive. For convenience it's OK to pass in a single string.98        If *regex* is ``True`` (default is ``False``) the given values are99        treated as regular expressions.100        """101        type_matched = self.ttype is ttype102        if not type_matched or values is None:103            return type_matched104 105        if isinstance(values, str):106            values = (values,)107 108        if regex:109            # TODO: Add test for regex with is_keyboard = false110            flag = re.IGNORECASE if self.is_keyword else 0111            values = (re.compile(v, flag) for v in values)112 113            for pattern in values:114                if pattern.search(self.normalized):115                    return True116            return False117 118        if self.is_keyword:119            values = (v.upper() for v in values)120 121        return self.normalized in values122 123    def within(self, group_cls):124        """Returns ``True`` if this token is within *group_cls*.125 126        Use this method for example to check if an identifier is within127        a function: ``t.within(sql.Function)``.128        """129        parent = self.parent130        while parent:131            if isinstance(parent, group_cls):132                return True133            parent = parent.parent134        return False135 136    def is_child_of(self, other):137        """Returns ``True`` if this token is a direct child of *other*."""138        return self.parent == other139 140    def has_ancestor(self, other):141        """Returns ``True`` if *other* is in this tokens ancestry."""142        parent = self.parent143        while parent:144            if parent == other:145                return True146            parent = parent.parent147        return False148 149 150class TokenList(Token):151    """A group of tokens.152 153    It has an additional instance attribute ``tokens`` which holds a154    list of child-tokens.155    """156 157    __slots__ = 'tokens'158 159    def __init__(self, tokens=None):160        self.tokens = tokens or []161        [setattr(token, 'parent', self) for token in self.tokens]162        super().__init__(None, str(self))163        self.is_group = True164 165    def __str__(self):166        return ''.join(token.value for token in self.flatten())167 168    # weird bug169    # def __len__(self):170    #     return len(self.tokens)171 172    def __iter__(self):173        return iter(self.tokens)174 175    def __getitem__(self, item):176        return self.tokens[item]177 178    def _get_repr_name(self):179        return type(self).__name__180 181    def _pprint_tree(self, max_depth=None, depth=0, f=None, _pre=''):182        """Pretty-print the object tree."""183        token_count = len(self.tokens)184        for idx, token in enumerate(self.tokens):185            cls = token._get_repr_name()186            value = token._get_repr_value()187 188            last = idx == (token_count - 1)189            pre = '`- ' if last else '|- '190 191            q = '"' if value.startswith("'") and value.endswith("'") else "'"192            print("{_pre}{pre}{idx} {cls} {q}{value}{q}"193                  .format(**locals()), file=f)194 195            if token.is_group and (max_depth is None or depth < max_depth):196                parent_pre = '   ' if last else '|  '197                token._pprint_tree(max_depth, depth + 1, f, _pre + parent_pre)198 199    def get_token_at_offset(self, offset):200        """Returns the token that is on position offset."""201        idx = 0202        for token in self.flatten():203            end = idx + len(token.value)204            if idx <= offset < end:205                return token206            idx = end207 208    def flatten(self):209        """Generator yielding ungrouped tokens.210 211        This method is recursively called for all child tokens.212        """213        try:214            for token in self.tokens:215                if token.is_group:216                    yield from token.flatten()217                else:218                    yield token219        except RecursionError as err:220            raise SQLParseError('Maximum recursion depth exceeded') from err221 222    def get_sublists(self):223        for token in self.tokens:224            if token.is_group:225                yield token226 227    @property228    def _groupable_tokens(self):229        return self.tokens230 231    def _token_matching(self, funcs, start=0, end=None, reverse=False):232        """next token that match functions"""233        if start is None:234            return None235 236        if not isinstance(funcs, (list, tuple)):237            funcs = (funcs,)238 239        if reverse:240            assert end is None241            indexes = range(start - 2, -1, -1)242        else:243            if end is None:244                end = len(self.tokens)245            indexes = range(start, end)246        for idx in indexes:247            token = self.tokens[idx]248            for func in funcs:249                if func(token):250                    return idx, token251        return None, None252 253    def token_first(self, skip_ws=True, skip_cm=False):254        """Returns the first child token.255 256        If *skip_ws* is ``True`` (the default), whitespace257        tokens are ignored.258 259        if *skip_cm* is ``True`` (default: ``False``), comments are260        ignored too.261        """262        # this on is inconsistent, using Comment instead of T.Comment...263        def matcher(tk):264            return not ((skip_ws and tk.is_whitespace)265                        or (skip_cm and imt(tk, t=T.Comment, i=Comment)))266        return self._token_matching(matcher)[1]267 268    def token_next_by(self, i=None, m=None, t=None, idx=-1, end=None):269        idx += 1270        return self._token_matching(lambda tk: imt(tk, i, m, t), idx, end)271 272    def token_not_matching(self, funcs, idx):273        funcs = (funcs,) if not isinstance(funcs, (list, tuple)) else funcs274        funcs = [lambda tk: not func(tk) for func in funcs]275        return self._token_matching(funcs, idx)276 277    def token_matching(self, funcs, idx):278        return self._token_matching(funcs, idx)[1]279 280    def token_prev(self, idx, skip_ws=True, skip_cm=False):281        """Returns the previous token relative to *idx*.282 283        If *skip_ws* is ``True`` (the default) whitespace tokens are ignored.284        If *skip_cm* is ``True`` comments are ignored.285        ``None`` is returned if there's no previous token.286        """287        return self.token_next(idx, skip_ws, skip_cm, _reverse=True)288 289    # TODO: May need to re-add default value to idx290    def token_next(self, idx, skip_ws=True, skip_cm=False, _reverse=False):291        """Returns the next token relative to *idx*.292 293        If *skip_ws* is ``True`` (the default) whitespace tokens are ignored.294        If *skip_cm* is ``True`` comments are ignored.295        ``None`` is returned if there's no next token.296        """297        if idx is None:298            return None, None299        idx += 1  # alot of code usage current pre-compensates for this300 301        def matcher(tk):302            return not ((skip_ws and tk.is_whitespace)303                        or (skip_cm and imt(tk, t=T.Comment, i=Comment)))304        return self._token_matching(matcher, idx, reverse=_reverse)305 306    def token_index(self, token, start=0):307        """Return list index of token."""308        start = start if isinstance(start, int) else self.token_index(start)309        return start + self.tokens[start:].index(token)310 311    def group_tokens(self, grp_cls, start, end, include_end=True,312                     extend=False):313        """Replace tokens by an instance of *grp_cls*."""314        start_idx = start315        start = self.tokens[start_idx]316 317        end_idx = end + include_end318 319        # will be needed later for new group_clauses320        # while skip_ws and tokens and tokens[-1].is_whitespace:321        #     tokens = tokens[:-1]322 323        if extend and isinstance(start, grp_cls):324            subtokens = self.tokens[start_idx + 1:end_idx]325 326            grp = start327            grp.tokens.extend(subtokens)328            del self.tokens[start_idx + 1:end_idx]329            grp.value = str(start)330        else:331            subtokens = self.tokens[start_idx:end_idx]332            grp = grp_cls(subtokens)333            self.tokens[start_idx:end_idx] = [grp]334            grp.parent = self335 336        for token in subtokens:337            token.parent = grp338 339        return grp340 341    def insert_before(self, where, token):342        """Inserts *token* before *where*."""343        if not isinstance(where, int):344            where = self.token_index(where)345        token.parent = self346        self.tokens.insert(where, token)347 348    def insert_after(self, where, token, skip_ws=True):349        """Inserts *token* after *where*."""350        if not isinstance(where, int):351            where = self.token_index(where)352        nidx, next_ = self.token_next(where, skip_ws=skip_ws)353        token.parent = self354        if next_ is None:355            self.tokens.append(token)356        else:357            self.tokens.insert(nidx, token)358 359    def has_alias(self):360        """Returns ``True`` if an alias is present."""361        return self.get_alias() is not None362 363    def get_alias(self):364        """Returns the alias for this identifier or ``None``."""365        return None366 367    def get_name(self):368        """Returns the name of this identifier.369 370        This is either it's alias or it's real name. The returned valued can371        be considered as the name under which the object corresponding to372        this identifier is known within the current statement.373        """374        return self.get_alias() or self.get_real_name()375 376    def get_real_name(self):377        """Returns the real name (object name) of this identifier."""378        return None379 380    def get_parent_name(self):381        """Return name of the parent object if any.382 383        A parent object is identified by the first occurring dot.384        """385        dot_idx, _ = self.token_next_by(m=(T.Punctuation, '.'))386        _, prev_ = self.token_prev(dot_idx)387        return remove_quotes(prev_.value) if prev_ is not None else None388 389    def _get_first_name(self, idx=None, reverse=False, keywords=False,390                        real_name=False):391        """Returns the name of the first token with a name"""392 393        tokens = self.tokens[idx:] if idx else self.tokens394        tokens = reversed(tokens) if reverse else tokens395        types = [T.Name, T.Wildcard, T.String.Symbol]396 397        if keywords:398            types.append(T.Keyword)399 400        for token in tokens:401            if token.ttype in types:402                return remove_quotes(token.value)403            elif isinstance(token, (Identifier, Function)):404                return token.get_real_name() if real_name else token.get_name()405 406 407class Statement(TokenList):408    """Represents a SQL statement."""409 410    def get_type(self):411        """Returns the type of a statement.412 413        The returned value is a string holding an upper-cased reprint of414        the first DML or DDL keyword. If the first token in this group415        isn't a DML or DDL keyword "UNKNOWN" is returned.416 417        Whitespaces and comments at the beginning of the statement418        are ignored.419        """420        token = self.token_first(skip_cm=True)421        if token is None:422            # An "empty" statement that either has not tokens at all423            # or only whitespace tokens.424            return 'UNKNOWN'425 426        elif token.ttype in (T.Keyword.DML, T.Keyword.DDL):427            return token.normalized428 429        elif token.ttype == T.Keyword.CTE:430            # The WITH keyword should be followed by either an Identifier or431            # an IdentifierList containing the CTE definitions;  the actual432            # DML keyword (e.g. SELECT, INSERT) will follow next.433            tidx = self.token_index(token)434            while tidx is not None:435                tidx, token = self.token_next(tidx, skip_ws=True)436                if isinstance(token, (Identifier, IdentifierList)):437                    tidx, token = self.token_next(tidx, skip_ws=True)438 439                    if token is not None \440                            and token.ttype == T.Keyword.DML:441                        return token.normalized442 443        # Hmm, probably invalid syntax, so return unknown.444        return 'UNKNOWN'445 446 447class Identifier(NameAliasMixin, TokenList):448    """Represents an identifier.449 450    Identifiers may have aliases or typecasts.451    """452 453    def is_wildcard(self):454        """Return ``True`` if this identifier contains a wildcard."""455        _, token = self.token_next_by(t=T.Wildcard)456        return token is not None457 458    def get_typecast(self):459        """Returns the typecast or ``None`` of this object as a string."""460        midx, marker = self.token_next_by(m=(T.Punctuation, '::'))461        nidx, next_ = self.token_next(midx, skip_ws=False)462        return next_.value if next_ else None463 464    def get_ordering(self):465        """Returns the ordering or ``None`` as uppercase string."""466        _, ordering = self.token_next_by(t=T.Keyword.Order)467        return ordering.normalized if ordering else None468 469    def get_array_indices(self):470        """Returns an iterator of index token lists"""471 472        for token in self.tokens:473            if isinstance(token, SquareBrackets):474                # Use [1:-1] index to discard the square brackets475                yield token.tokens[1:-1]476 477 478class IdentifierList(TokenList):479    """A list of :class:`~sqlparse.sql.Identifier`\'s."""480 481    def get_identifiers(self):482        """Returns the identifiers.483 484        Whitespaces and punctuations are not included in this generator.485        """486        for token in self.tokens:487            if not (token.is_whitespace or token.match(T.Punctuation, ',')):488                yield token489 490 491class TypedLiteral(TokenList):492    """A typed literal, such as "date '2001-09-28'" or "interval '2 hours'"."""493    M_OPEN = [(T.Name.Builtin, None), (T.Keyword, "TIMESTAMP")]494    M_CLOSE = T.String.Single, None495    M_EXTEND = T.Keyword, ("DAY", "HOUR", "MINUTE", "MONTH", "SECOND", "YEAR")496 497 498class Parenthesis(TokenList):499    """Tokens between parenthesis."""500    M_OPEN = T.Punctuation, '('501    M_CLOSE = T.Punctuation, ')'502 503    @property504    def _groupable_tokens(self):505        return self.tokens[1:-1]506 507 508class SquareBrackets(TokenList):509    """Tokens between square brackets"""510    M_OPEN = T.Punctuation, '['511    M_CLOSE = T.Punctuation, ']'512 513    @property514    def _groupable_tokens(self):515        return self.tokens[1:-1]516 517 518class Assignment(TokenList):519    """An assignment like 'var := val;'"""520 521 522class If(TokenList):523    """An 'if' clause with possible 'else if' or 'else' parts."""524    M_OPEN = T.Keyword, 'IF'525    M_CLOSE = T.Keyword, 'END IF'526 527 528class For(TokenList):529    """A 'FOR' loop."""530    M_OPEN = T.Keyword, ('FOR', 'FOREACH')531    M_CLOSE = T.Keyword, 'END LOOP'532 533 534class Comparison(TokenList):535    """A comparison used for example in WHERE clauses."""536 537    @property538    def left(self):539        return self.tokens[0]540 541    @property542    def right(self):543        return self.tokens[-1]544 545 546class Comment(TokenList):547    """A comment."""548 549    def is_multiline(self):550        return self.tokens and self.tokens[0].ttype == T.Comment.Multiline551 552 553class Where(TokenList):554    """A WHERE clause."""555    M_OPEN = T.Keyword, 'WHERE'556    M_CLOSE = T.Keyword, (557        'ORDER BY', 'GROUP BY', 'LIMIT', 'UNION', 'UNION ALL', 'EXCEPT',558        'HAVING', 'RETURNING', 'INTO')559 560 561class Over(TokenList):562    """An OVER clause."""563    M_OPEN = T.Keyword, 'OVER'564 565 566class Having(TokenList):567    """A HAVING clause."""568    M_OPEN = T.Keyword, 'HAVING'569    M_CLOSE = T.Keyword, ('ORDER BY', 'LIMIT')570 571 572class Case(TokenList):573    """A CASE statement with one or more WHEN and possibly an ELSE part."""574    M_OPEN = T.Keyword, 'CASE'575    M_CLOSE = T.Keyword, 'END'576 577    def get_cases(self, skip_ws=False):578        """Returns a list of 2-tuples (condition, value).579 580        If an ELSE exists condition is None.581        """582        CONDITION = 1583        VALUE = 2584 585        ret = []586        mode = CONDITION587 588        for token in self.tokens:589            # Set mode from the current statement590            if token.match(T.Keyword, 'CASE'):591                continue592 593            elif skip_ws and token.ttype in T.Whitespace:594                continue595 596            elif token.match(T.Keyword, 'WHEN'):597                ret.append(([], []))598                mode = CONDITION599 600            elif token.match(T.Keyword, 'THEN'):601                mode = VALUE602 603            elif token.match(T.Keyword, 'ELSE'):604                ret.append((None, []))605                mode = VALUE606 607            elif token.match(T.Keyword, 'END'):608                mode = None609 610            # First condition without preceding WHEN611            if mode and not ret:612                ret.append(([], []))613 614            # Append token depending of the current mode615            if mode == CONDITION:616                ret[-1][0].append(token)617 618            elif mode == VALUE:619                ret[-1][1].append(token)620 621        # Return cases list622        return ret623 624 625class Function(NameAliasMixin, TokenList):626    """A function or procedure call."""627 628    def get_parameters(self):629        """Return a list of parameters."""630        parenthesis = self.token_next_by(i=Parenthesis)[1]631        result = []632        for token in parenthesis.tokens:633            if isinstance(token, IdentifierList):634                return token.get_identifiers()635            elif imt(token, i=(Function, Identifier, TypedLiteral),636                     t=T.Literal):637                result.append(token)638        return result639 640    def get_window(self):641        """Return the window if it exists."""642        over_clause = self.token_next_by(i=Over)643        if not over_clause:644            return None645        return over_clause[1].tokens[-1]646 647 648class Begin(TokenList):649    """A BEGIN/END block."""650    M_OPEN = T.Keyword, 'BEGIN'651    M_CLOSE = T.Keyword, 'END'652 653 654class Operation(TokenList):655    """Grouping of operations"""656 657 658class Values(TokenList):659    """Grouping of values"""660 661 662class Command(TokenList):663    """Grouping of CLI commands."""664 
codekingpro/portable-devtools · Team Ai