-
Notifications
You must be signed in to change notification settings - Fork 21
/
identifier.py
93 lines (73 loc) · 3.03 KB
/
identifier.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import re
from copy import copy, deepcopy
from mindsdb_sql.parser.ast.base import ASTNode
from mindsdb_sql.parser.utils import indent
from mindsdb_sql.parser.ast.select import Star
no_wrap_identifier_regex = re.compile(r'[a-zA-Z_][a-zA-Z_0-9]*')
path_str_parts_regex = re.compile(r'(?:(?:(`[^`]+`))|([^.]+))')
def path_str_to_parts(path_str: str):
match = re.finditer(path_str_parts_regex, path_str)
parts = [x[0].strip('`') for x in match]
return parts
RESERVED_KEYWORDS = {
'PERSIST', 'IF', 'EXISTS', 'NULLS', 'FIRST', 'LAST',
'ORDER', 'BY', 'GROUP', 'PARTITION'
}
def get_reserved_words():
from mindsdb_sql.parser.lexer import SQLLexer
from mindsdb_sql.parser.dialects.mindsdb.lexer import MindsDBLexer
reserved = RESERVED_KEYWORDS
for word in SQLLexer.tokens | MindsDBLexer.tokens:
if '_' not in word:
# exclude combinations
reserved.add(word)
return reserved
class Identifier(ASTNode):
def __init__(self, path_str=None, parts=None, *args, **kwargs):
super().__init__(*args, **kwargs)
assert path_str or parts, "Either path_str or parts must be provided for an Identifier"
assert not (path_str and parts), "Provide either path_str or parts, but not both"
if isinstance(path_str, Star) and not parts:
parts = [Star()]
if path_str and not parts:
parts = path_str_to_parts(path_str)
assert isinstance(parts, list)
self.parts = parts
@classmethod
def from_path_str(self, value, *args, **kwargs):
parts = path_str_to_parts(value)
return Identifier(parts=parts, *args, **kwargs)
def parts_to_str(self):
out_parts = []
reserved_words = get_reserved_words()
for part in self.parts:
if isinstance(part, Star):
part = str(part)
else:
if (
not no_wrap_identifier_regex.fullmatch(part)
or
part.upper() in reserved_words
):
part = f'`{part}`'
out_parts.append(part)
return '.'.join(out_parts)
def to_tree(self, *args, level=0, **kwargs):
alias_str = f', alias={self.alias.to_tree()}' if self.alias else ''
return indent(level) + f'Identifier(parts={[str(i) for i in self.parts]}{alias_str})'
def get_string(self, *args, **kwargs):
return self.parts_to_str()
def __copy__(self):
identifier = Identifier(parts=copy(self.parts))
identifier.alias = deepcopy(self.alias)
identifier.parentheses = self.parentheses
if hasattr(self, 'sub_select'):
identifier.sub_select = deepcopy(self.sub_select)
return identifier
def __deepcopy__(self, memo):
identifier = Identifier(parts=copy(self.parts))
identifier.alias = deepcopy(self.alias)
identifier.parentheses = self.parentheses
if hasattr(self, 'sub_select'):
identifier.sub_select = deepcopy(self.sub_select)
return identifier