pytrilogy 0.3.142__cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (200) hide show
  1. LICENSE.md +19 -0
  2. _preql_import_resolver/__init__.py +5 -0
  3. _preql_import_resolver/_preql_import_resolver.cpython-313-x86_64-linux-gnu.so +0 -0
  4. pytrilogy-0.3.142.dist-info/METADATA +555 -0
  5. pytrilogy-0.3.142.dist-info/RECORD +200 -0
  6. pytrilogy-0.3.142.dist-info/WHEEL +5 -0
  7. pytrilogy-0.3.142.dist-info/entry_points.txt +2 -0
  8. pytrilogy-0.3.142.dist-info/licenses/LICENSE.md +19 -0
  9. trilogy/__init__.py +16 -0
  10. trilogy/ai/README.md +10 -0
  11. trilogy/ai/__init__.py +19 -0
  12. trilogy/ai/constants.py +92 -0
  13. trilogy/ai/conversation.py +107 -0
  14. trilogy/ai/enums.py +7 -0
  15. trilogy/ai/execute.py +50 -0
  16. trilogy/ai/models.py +34 -0
  17. trilogy/ai/prompts.py +100 -0
  18. trilogy/ai/providers/__init__.py +0 -0
  19. trilogy/ai/providers/anthropic.py +106 -0
  20. trilogy/ai/providers/base.py +24 -0
  21. trilogy/ai/providers/google.py +146 -0
  22. trilogy/ai/providers/openai.py +89 -0
  23. trilogy/ai/providers/utils.py +68 -0
  24. trilogy/authoring/README.md +3 -0
  25. trilogy/authoring/__init__.py +148 -0
  26. trilogy/constants.py +113 -0
  27. trilogy/core/README.md +52 -0
  28. trilogy/core/__init__.py +0 -0
  29. trilogy/core/constants.py +6 -0
  30. trilogy/core/enums.py +443 -0
  31. trilogy/core/env_processor.py +120 -0
  32. trilogy/core/environment_helpers.py +320 -0
  33. trilogy/core/ergonomics.py +193 -0
  34. trilogy/core/exceptions.py +123 -0
  35. trilogy/core/functions.py +1227 -0
  36. trilogy/core/graph_models.py +139 -0
  37. trilogy/core/internal.py +85 -0
  38. trilogy/core/models/__init__.py +0 -0
  39. trilogy/core/models/author.py +2669 -0
  40. trilogy/core/models/build.py +2521 -0
  41. trilogy/core/models/build_environment.py +180 -0
  42. trilogy/core/models/core.py +501 -0
  43. trilogy/core/models/datasource.py +322 -0
  44. trilogy/core/models/environment.py +751 -0
  45. trilogy/core/models/execute.py +1177 -0
  46. trilogy/core/optimization.py +251 -0
  47. trilogy/core/optimizations/__init__.py +12 -0
  48. trilogy/core/optimizations/base_optimization.py +17 -0
  49. trilogy/core/optimizations/hide_unused_concept.py +47 -0
  50. trilogy/core/optimizations/inline_datasource.py +102 -0
  51. trilogy/core/optimizations/predicate_pushdown.py +245 -0
  52. trilogy/core/processing/README.md +94 -0
  53. trilogy/core/processing/READMEv2.md +121 -0
  54. trilogy/core/processing/VIRTUAL_UNNEST.md +30 -0
  55. trilogy/core/processing/__init__.py +0 -0
  56. trilogy/core/processing/concept_strategies_v3.py +508 -0
  57. trilogy/core/processing/constants.py +15 -0
  58. trilogy/core/processing/discovery_node_factory.py +451 -0
  59. trilogy/core/processing/discovery_utility.py +548 -0
  60. trilogy/core/processing/discovery_validation.py +167 -0
  61. trilogy/core/processing/graph_utils.py +43 -0
  62. trilogy/core/processing/node_generators/README.md +9 -0
  63. trilogy/core/processing/node_generators/__init__.py +31 -0
  64. trilogy/core/processing/node_generators/basic_node.py +160 -0
  65. trilogy/core/processing/node_generators/common.py +268 -0
  66. trilogy/core/processing/node_generators/constant_node.py +38 -0
  67. trilogy/core/processing/node_generators/filter_node.py +315 -0
  68. trilogy/core/processing/node_generators/group_node.py +213 -0
  69. trilogy/core/processing/node_generators/group_to_node.py +117 -0
  70. trilogy/core/processing/node_generators/multiselect_node.py +205 -0
  71. trilogy/core/processing/node_generators/node_merge_node.py +653 -0
  72. trilogy/core/processing/node_generators/recursive_node.py +88 -0
  73. trilogy/core/processing/node_generators/rowset_node.py +165 -0
  74. trilogy/core/processing/node_generators/select_helpers/__init__.py +0 -0
  75. trilogy/core/processing/node_generators/select_helpers/datasource_injection.py +261 -0
  76. trilogy/core/processing/node_generators/select_merge_node.py +748 -0
  77. trilogy/core/processing/node_generators/select_node.py +95 -0
  78. trilogy/core/processing/node_generators/synonym_node.py +98 -0
  79. trilogy/core/processing/node_generators/union_node.py +91 -0
  80. trilogy/core/processing/node_generators/unnest_node.py +182 -0
  81. trilogy/core/processing/node_generators/window_node.py +201 -0
  82. trilogy/core/processing/nodes/README.md +28 -0
  83. trilogy/core/processing/nodes/__init__.py +179 -0
  84. trilogy/core/processing/nodes/base_node.py +519 -0
  85. trilogy/core/processing/nodes/filter_node.py +75 -0
  86. trilogy/core/processing/nodes/group_node.py +194 -0
  87. trilogy/core/processing/nodes/merge_node.py +420 -0
  88. trilogy/core/processing/nodes/recursive_node.py +46 -0
  89. trilogy/core/processing/nodes/select_node_v2.py +242 -0
  90. trilogy/core/processing/nodes/union_node.py +53 -0
  91. trilogy/core/processing/nodes/unnest_node.py +62 -0
  92. trilogy/core/processing/nodes/window_node.py +56 -0
  93. trilogy/core/processing/utility.py +823 -0
  94. trilogy/core/query_processor.py +596 -0
  95. trilogy/core/statements/README.md +35 -0
  96. trilogy/core/statements/__init__.py +0 -0
  97. trilogy/core/statements/author.py +536 -0
  98. trilogy/core/statements/build.py +0 -0
  99. trilogy/core/statements/common.py +20 -0
  100. trilogy/core/statements/execute.py +155 -0
  101. trilogy/core/table_processor.py +66 -0
  102. trilogy/core/utility.py +8 -0
  103. trilogy/core/validation/README.md +46 -0
  104. trilogy/core/validation/__init__.py +0 -0
  105. trilogy/core/validation/common.py +161 -0
  106. trilogy/core/validation/concept.py +146 -0
  107. trilogy/core/validation/datasource.py +227 -0
  108. trilogy/core/validation/environment.py +73 -0
  109. trilogy/core/validation/fix.py +256 -0
  110. trilogy/dialect/__init__.py +32 -0
  111. trilogy/dialect/base.py +1392 -0
  112. trilogy/dialect/bigquery.py +308 -0
  113. trilogy/dialect/common.py +147 -0
  114. trilogy/dialect/config.py +144 -0
  115. trilogy/dialect/dataframe.py +50 -0
  116. trilogy/dialect/duckdb.py +231 -0
  117. trilogy/dialect/enums.py +147 -0
  118. trilogy/dialect/metadata.py +173 -0
  119. trilogy/dialect/mock.py +190 -0
  120. trilogy/dialect/postgres.py +117 -0
  121. trilogy/dialect/presto.py +110 -0
  122. trilogy/dialect/results.py +89 -0
  123. trilogy/dialect/snowflake.py +129 -0
  124. trilogy/dialect/sql_server.py +137 -0
  125. trilogy/engine.py +48 -0
  126. trilogy/execution/config.py +75 -0
  127. trilogy/executor.py +568 -0
  128. trilogy/hooks/__init__.py +4 -0
  129. trilogy/hooks/base_hook.py +40 -0
  130. trilogy/hooks/graph_hook.py +139 -0
  131. trilogy/hooks/query_debugger.py +166 -0
  132. trilogy/metadata/__init__.py +0 -0
  133. trilogy/parser.py +10 -0
  134. trilogy/parsing/README.md +21 -0
  135. trilogy/parsing/__init__.py +0 -0
  136. trilogy/parsing/common.py +1069 -0
  137. trilogy/parsing/config.py +5 -0
  138. trilogy/parsing/exceptions.py +8 -0
  139. trilogy/parsing/helpers.py +1 -0
  140. trilogy/parsing/parse_engine.py +2813 -0
  141. trilogy/parsing/render.py +769 -0
  142. trilogy/parsing/trilogy.lark +540 -0
  143. trilogy/py.typed +0 -0
  144. trilogy/render.py +42 -0
  145. trilogy/scripts/README.md +9 -0
  146. trilogy/scripts/__init__.py +0 -0
  147. trilogy/scripts/agent.py +41 -0
  148. trilogy/scripts/agent_info.py +303 -0
  149. trilogy/scripts/common.py +355 -0
  150. trilogy/scripts/dependency/Cargo.lock +617 -0
  151. trilogy/scripts/dependency/Cargo.toml +39 -0
  152. trilogy/scripts/dependency/README.md +131 -0
  153. trilogy/scripts/dependency/build.sh +25 -0
  154. trilogy/scripts/dependency/src/directory_resolver.rs +177 -0
  155. trilogy/scripts/dependency/src/lib.rs +16 -0
  156. trilogy/scripts/dependency/src/main.rs +770 -0
  157. trilogy/scripts/dependency/src/parser.rs +435 -0
  158. trilogy/scripts/dependency/src/preql.pest +208 -0
  159. trilogy/scripts/dependency/src/python_bindings.rs +303 -0
  160. trilogy/scripts/dependency/src/resolver.rs +716 -0
  161. trilogy/scripts/dependency/tests/base.preql +3 -0
  162. trilogy/scripts/dependency/tests/cli_integration.rs +377 -0
  163. trilogy/scripts/dependency/tests/customer.preql +6 -0
  164. trilogy/scripts/dependency/tests/main.preql +9 -0
  165. trilogy/scripts/dependency/tests/orders.preql +7 -0
  166. trilogy/scripts/dependency/tests/test_data/base.preql +9 -0
  167. trilogy/scripts/dependency/tests/test_data/consumer.preql +1 -0
  168. trilogy/scripts/dependency.py +323 -0
  169. trilogy/scripts/display.py +512 -0
  170. trilogy/scripts/environment.py +46 -0
  171. trilogy/scripts/fmt.py +32 -0
  172. trilogy/scripts/ingest.py +471 -0
  173. trilogy/scripts/ingest_helpers/__init__.py +1 -0
  174. trilogy/scripts/ingest_helpers/foreign_keys.py +123 -0
  175. trilogy/scripts/ingest_helpers/formatting.py +93 -0
  176. trilogy/scripts/ingest_helpers/typing.py +161 -0
  177. trilogy/scripts/init.py +105 -0
  178. trilogy/scripts/parallel_execution.py +713 -0
  179. trilogy/scripts/plan.py +189 -0
  180. trilogy/scripts/run.py +63 -0
  181. trilogy/scripts/serve.py +140 -0
  182. trilogy/scripts/serve_helpers/__init__.py +41 -0
  183. trilogy/scripts/serve_helpers/file_discovery.py +142 -0
  184. trilogy/scripts/serve_helpers/index_generation.py +206 -0
  185. trilogy/scripts/serve_helpers/models.py +38 -0
  186. trilogy/scripts/single_execution.py +131 -0
  187. trilogy/scripts/testing.py +119 -0
  188. trilogy/scripts/trilogy.py +68 -0
  189. trilogy/std/__init__.py +0 -0
  190. trilogy/std/color.preql +3 -0
  191. trilogy/std/date.preql +13 -0
  192. trilogy/std/display.preql +18 -0
  193. trilogy/std/geography.preql +22 -0
  194. trilogy/std/metric.preql +15 -0
  195. trilogy/std/money.preql +67 -0
  196. trilogy/std/net.preql +14 -0
  197. trilogy/std/ranking.preql +7 -0
  198. trilogy/std/report.preql +5 -0
  199. trilogy/std/semantic.preql +6 -0
  200. trilogy/utility.py +34 -0
@@ -0,0 +1,2669 @@
1
+ from __future__ import annotations
2
+
3
+ import hashlib
4
+ from abc import ABC
5
+ from datetime import date, datetime
6
+ from functools import cached_property
7
+ from typing import (
8
+ TYPE_CHECKING,
9
+ Any,
10
+ Iterable,
11
+ List,
12
+ Optional,
13
+ Self,
14
+ Sequence,
15
+ Set,
16
+ Tuple,
17
+ Type,
18
+ Union,
19
+ )
20
+
21
+ from pydantic import (
22
+ BaseModel,
23
+ ConfigDict,
24
+ Field,
25
+ ValidationInfo,
26
+ computed_field,
27
+ field_validator,
28
+ model_validator,
29
+ )
30
+
31
+ from trilogy.constants import DEFAULT_NAMESPACE, MagicConstants
32
+ from trilogy.core.constants import ALL_ROWS_CONCEPT
33
+ from trilogy.core.enums import (
34
+ BooleanOperator,
35
+ ComparisonOperator,
36
+ ConceptSource,
37
+ DatePart,
38
+ Derivation,
39
+ FunctionClass,
40
+ FunctionType,
41
+ Granularity,
42
+ InfiniteFunctionArgs,
43
+ Modifier,
44
+ Ordering,
45
+ Purpose,
46
+ WindowOrder,
47
+ WindowType,
48
+ )
49
+ from trilogy.core.models.core import (
50
+ Addressable,
51
+ ArrayType,
52
+ DataType,
53
+ DataTyped,
54
+ ListWrapper,
55
+ MapType,
56
+ MapWrapper,
57
+ NumericType,
58
+ StructType,
59
+ TraitDataType,
60
+ TupleWrapper,
61
+ arg_to_datatype,
62
+ is_compatible_datatype,
63
+ )
64
+ from trilogy.utility import unique
65
+
66
+ # TODO: refactor to avoid these
67
+ if TYPE_CHECKING:
68
+ from trilogy.core.models.environment import Environment
69
+
70
+
71
+ class Namespaced(ABC):
72
+ def with_namespace(self, namespace: str):
73
+ raise NotImplementedError
74
+
75
+
76
+ class Mergeable(ABC):
77
+ def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
78
+ raise NotImplementedError
79
+
80
+ def with_reference_replacement(self, source: str, target: Expr | ArgBinding):
81
+ raise NotImplementedError(type(self))
82
+
83
+
84
+ class ConceptArgs(ABC):
85
+ @property
86
+ def concept_arguments(self) -> Sequence["ConceptRef"]:
87
+ raise NotImplementedError
88
+
89
+ @property
90
+ def existence_arguments(self) -> Sequence[tuple["ConceptRef", ...]]:
91
+ return []
92
+
93
+ @property
94
+ def row_arguments(self) -> Sequence["ConceptRef"]:
95
+ return self.concept_arguments
96
+
97
+
98
+ class HasUUID(ABC):
99
+ @property
100
+ def uuid(self) -> str:
101
+ return hashlib.md5(str(self).encode()).hexdigest()
102
+
103
+
104
+ class ConceptRef(Addressable, Namespaced, DataTyped, Mergeable, BaseModel):
105
+ address: str
106
+ datatype: (
107
+ DataType | TraitDataType | ArrayType | StructType | MapType | NumericType
108
+ ) = DataType.UNKNOWN
109
+ metadata: Optional["Metadata"] = None
110
+
111
+ @property
112
+ def reference(self):
113
+ return self
114
+
115
+ @property
116
+ def line_no(self) -> int | None:
117
+ if self.metadata:
118
+ return self.metadata.line_number
119
+ return None
120
+
121
+ def __repr__(self):
122
+ return f"ref:{self.address}"
123
+
124
+ def __str__(self):
125
+ return self.__repr__()
126
+
127
+ def __eq__(self, other):
128
+ if isinstance(other, Concept):
129
+ return self.address == other.address
130
+ elif isinstance(other, str):
131
+ return self.address == other
132
+ elif isinstance(other, ConceptRef):
133
+ return self.address == other.address
134
+ return False
135
+
136
+ @property
137
+ def namespace(self):
138
+ return self.address.rsplit(".", 1)[0]
139
+
140
+ @property
141
+ def name(self):
142
+ return self.address.rsplit(".", 1)[1]
143
+
144
+ @property
145
+ def output_datatype(self):
146
+ return self.datatype
147
+
148
+ def with_merge(
149
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
150
+ ) -> ConceptRef:
151
+ if self.address == source.address:
152
+ return ConceptRef.model_construct(
153
+ address=target.address, datatype=target.datatype, metadata=self.metadata
154
+ )
155
+ return self
156
+
157
+ def with_namespace(self, namespace: str):
158
+ return ConceptRef.model_construct(
159
+ address=address_with_namespace(self.address, namespace),
160
+ datatype=self.datatype,
161
+ metadata=self.metadata,
162
+ )
163
+
164
+ def with_reference_replacement(self, source: str, target: Expr | ArgBinding):
165
+ # a reference might be to an attribute of a struct that is bound late
166
+ # if the replacement is a parent in the access path; replace reference
167
+ # with an attribute access call
168
+ candidates = [f"{DEFAULT_NAMESPACE}.{self.address}", self.address]
169
+ for candidate in candidates:
170
+ if candidate == source:
171
+ return target
172
+ if not candidate.startswith(f"{source}."):
173
+ continue
174
+ attribute = self.address.rsplit(".", 1)[1]
175
+ dtype = arg_to_datatype(target)
176
+ if not isinstance(dtype, StructType):
177
+ continue
178
+ output_type = dtype.field_types.get(attribute, DataType.UNKNOWN)
179
+ return Function(
180
+ arguments=[target, self.address.rsplit(".", 1)[1]],
181
+ operator=FunctionType.ATTR_ACCESS,
182
+ arg_count=2,
183
+ output_datatype=output_type,
184
+ output_purpose=Purpose.PROPERTY,
185
+ )
186
+ return self
187
+
188
+
189
+ class UndefinedConcept(ConceptRef):
190
+ pass
191
+
192
+ @property
193
+ def reference(self):
194
+ return self
195
+
196
+ @property
197
+ def purpose(self):
198
+ return Purpose.UNKNOWN
199
+
200
+
201
+ def address_with_namespace(address: str, namespace: str) -> str:
202
+ existing_ns = address.split(".", 1)[0]
203
+ if "." in address:
204
+ existing_name = address.split(".", 1)[1]
205
+ else:
206
+ existing_name = address
207
+ if existing_name == ALL_ROWS_CONCEPT:
208
+ return address
209
+ if existing_ns == DEFAULT_NAMESPACE:
210
+ return f"{namespace}.{existing_name}"
211
+ return f"{namespace}.{address}"
212
+
213
+
214
+ class Parenthetical(
215
+ DataTyped,
216
+ ConceptArgs,
217
+ Mergeable,
218
+ Namespaced,
219
+ BaseModel,
220
+ ):
221
+ content: "Expr"
222
+
223
+ @field_validator("content", mode="before")
224
+ @classmethod
225
+ def content_validator(cls, v, info: ValidationInfo):
226
+ if isinstance(v, Concept):
227
+ return v.reference
228
+ return v
229
+
230
+ def __add__(self, other) -> Union["Parenthetical", "Conditional"]:
231
+ if other is None:
232
+ return self
233
+ elif isinstance(other, (Comparison, Conditional, Parenthetical)):
234
+ return Conditional(left=self, right=other, operator=BooleanOperator.AND)
235
+ raise ValueError(f"Cannot add {self.__class__} and {type(other)}")
236
+
237
+ def __str__(self):
238
+ return self.__repr__()
239
+
240
+ def __repr__(self):
241
+ return f"({str(self.content)})"
242
+
243
+ def with_namespace(self, namespace: str) -> Parenthetical:
244
+ return Parenthetical.model_construct(
245
+ content=(
246
+ self.content.with_namespace(namespace)
247
+ if isinstance(self.content, Namespaced)
248
+ else self.content
249
+ )
250
+ )
251
+
252
+ def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
253
+ return Parenthetical.model_construct(
254
+ content=(
255
+ self.content.with_merge(source, target, modifiers)
256
+ if isinstance(self.content, Mergeable)
257
+ else self.content
258
+ )
259
+ )
260
+
261
+ def with_reference_replacement(self, source, target):
262
+ return Parenthetical.model_construct(
263
+ content=(
264
+ self.content.with_reference_replacement(source, target)
265
+ if isinstance(self.content, Mergeable)
266
+ else self.content
267
+ )
268
+ )
269
+
270
+ @property
271
+ def concept_arguments(self) -> Sequence[ConceptRef]:
272
+ base: List[ConceptRef] = []
273
+ x = self.content
274
+ if isinstance(x, ConceptRef):
275
+ base += [x]
276
+ elif isinstance(x, ConceptArgs):
277
+ base += x.concept_arguments
278
+ return base
279
+
280
+ @property
281
+ def row_arguments(self) -> Sequence[ConceptRef]:
282
+ if isinstance(self.content, ConceptArgs):
283
+ return self.content.row_arguments
284
+ return self.concept_arguments
285
+
286
+ @property
287
+ def existence_arguments(self) -> Sequence[tuple["ConceptRef", ...]]:
288
+ if isinstance(self.content, ConceptArgs):
289
+ return self.content.existence_arguments
290
+ return []
291
+
292
+ @property
293
+ def output_datatype(self):
294
+ return arg_to_datatype(self.content)
295
+
296
+
297
+ class Conditional(Mergeable, ConceptArgs, Namespaced, DataTyped, BaseModel):
298
+ left: Expr
299
+ right: Expr
300
+ operator: BooleanOperator
301
+
302
+ @field_validator("left", mode="before")
303
+ @classmethod
304
+ def left_validator(cls, v, info: ValidationInfo):
305
+ if isinstance(v, Concept):
306
+ return v.reference
307
+ return v
308
+
309
+ @field_validator("right", mode="before")
310
+ @classmethod
311
+ def right_validator(cls, v, info: ValidationInfo):
312
+ if isinstance(v, Concept):
313
+ return v.reference
314
+ return v
315
+
316
+ def __add__(self, other) -> "Conditional":
317
+ if other is None:
318
+ return self
319
+ elif str(other) == str(self):
320
+ return self
321
+ elif isinstance(other, (Comparison, Conditional, Parenthetical)):
322
+ return Conditional.model_construct(
323
+ left=self, right=other, operator=BooleanOperator.AND
324
+ )
325
+ raise ValueError(f"Cannot add {self.__class__} and {type(other)}")
326
+
327
+ def __str__(self):
328
+ return self.__repr__()
329
+
330
+ def __repr__(self):
331
+ return f"{str(self.left)} {self.operator.value} {str(self.right)}"
332
+
333
+ def __eq__(self, other):
334
+ if not isinstance(other, Conditional):
335
+ return False
336
+ return (
337
+ self.left == other.left
338
+ and self.right == other.right
339
+ and self.operator == other.operator
340
+ )
341
+
342
+ def with_namespace(self, namespace: str) -> "Conditional":
343
+ return Conditional.model_construct(
344
+ left=(
345
+ self.left.with_namespace(namespace)
346
+ if isinstance(self.left, Namespaced)
347
+ else self.left
348
+ ),
349
+ right=(
350
+ self.right.with_namespace(namespace)
351
+ if isinstance(self.right, Namespaced)
352
+ else self.right
353
+ ),
354
+ operator=self.operator,
355
+ )
356
+
357
+ def with_merge(
358
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
359
+ ) -> "Conditional":
360
+ return Conditional.model_construct(
361
+ left=(
362
+ self.left.with_merge(source, target, modifiers)
363
+ if isinstance(self.left, Mergeable)
364
+ else self.left
365
+ ),
366
+ right=(
367
+ self.right.with_merge(source, target, modifiers)
368
+ if isinstance(self.right, Mergeable)
369
+ else self.right
370
+ ),
371
+ operator=self.operator,
372
+ )
373
+
374
+ def with_reference_replacement(self, source, target):
375
+ return self.__class__.model_construct(
376
+ left=(
377
+ self.left.with_reference_replacement(source, target)
378
+ if isinstance(self.left, Mergeable)
379
+ else self.left
380
+ ),
381
+ right=(
382
+ self.right.with_reference_replacement(source, target)
383
+ if isinstance(self.right, Mergeable)
384
+ else self.right
385
+ ),
386
+ operator=self.operator,
387
+ )
388
+
389
+ @property
390
+ def concept_arguments(self) -> Sequence[ConceptRef]:
391
+ output = []
392
+ output += get_concept_arguments(self.left)
393
+ output += get_concept_arguments(self.right)
394
+ return output
395
+
396
+ @property
397
+ def row_arguments(self) -> Sequence[ConceptRef]:
398
+ output = []
399
+ output += get_concept_row_arguments(self.left)
400
+ output += get_concept_row_arguments(self.right)
401
+ return output
402
+
403
+ @property
404
+ def existence_arguments(self) -> Sequence[tuple[ConceptRef, ...]]:
405
+ output: list[tuple[ConceptRef, ...]] = []
406
+ if isinstance(self.left, ConceptArgs):
407
+ output += self.left.existence_arguments
408
+ if isinstance(self.right, ConceptArgs):
409
+ output += self.right.existence_arguments
410
+ return output
411
+
412
+ @property
413
+ def output_datatype(self):
414
+ # a conditional is always a boolean
415
+ return DataType.BOOL
416
+
417
+ def decompose(self):
418
+ chunks = []
419
+ if self.operator == BooleanOperator.AND:
420
+ for val in [self.left, self.right]:
421
+ if isinstance(val, Conditional):
422
+ chunks.extend(val.decompose())
423
+ else:
424
+ chunks.append(val)
425
+ else:
426
+ chunks.append(self)
427
+ return chunks
428
+
429
+
430
+ class WhereClause(Mergeable, ConceptArgs, Namespaced, BaseModel):
431
+ conditional: Union[SubselectComparison, Comparison, Conditional, Parenthetical]
432
+
433
+ def __repr__(self):
434
+ return str(self.conditional)
435
+
436
+ def __str__(self):
437
+ return self.__repr__()
438
+
439
+ @property
440
+ def concept_arguments(self) -> Sequence[ConceptRef]:
441
+ return self.conditional.concept_arguments
442
+
443
+ @property
444
+ def row_arguments(self) -> Sequence[ConceptRef]:
445
+ return self.conditional.row_arguments
446
+
447
+ @property
448
+ def existence_arguments(self) -> Sequence[tuple["ConceptRef", ...]]:
449
+ return self.conditional.existence_arguments
450
+
451
+ def with_merge(
452
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
453
+ ) -> Self:
454
+ return self.__class__.model_construct(
455
+ conditional=self.conditional.with_merge(source, target, modifiers)
456
+ )
457
+
458
+ def with_namespace(self, namespace: str) -> Self:
459
+ return self.__class__.model_construct(
460
+ conditional=self.conditional.with_namespace(namespace)
461
+ )
462
+
463
+ def with_reference_replacement(self, source, target):
464
+ return self.__class__.model_construct(
465
+ conditional=self.conditional.with_reference_replacement(source, target)
466
+ )
467
+
468
+
469
+ class HavingClause(WhereClause):
470
+ pass
471
+
472
+
473
+ class Grain(Namespaced, BaseModel):
474
+ components: set[str] = Field(default_factory=set)
475
+ where_clause: Optional["WhereClause"] = None
476
+ _str: str | None = None
477
+ _abstract: bool = False
478
+
479
+ def without_condition(self):
480
+ return Grain(components=self.components)
481
+
482
+ def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
483
+ new_components = set()
484
+ for c in self.components:
485
+ if c == source.address:
486
+ new_components.add(target.address)
487
+ else:
488
+ new_components.add(c)
489
+ return Grain.model_construct(components=new_components)
490
+
491
+ @classmethod
492
+ def from_concepts(
493
+ cls,
494
+ concepts: Iterable[Concept | ConceptRef | str],
495
+ environment: Environment | None = None,
496
+ where_clause: WhereClause | None = None,
497
+ local_concepts: dict[str, Concept] | None = None,
498
+ ) -> Grain:
499
+ from trilogy.parsing.common import concepts_to_grain_concepts
500
+
501
+ x = Grain.model_construct(
502
+ components=concepts_to_grain_concepts(
503
+ concepts, environment=environment, local_concepts=local_concepts
504
+ ),
505
+ where_clause=where_clause,
506
+ )
507
+
508
+ return x
509
+
510
+ def with_namespace(self, namespace: str) -> "Grain":
511
+ return Grain.model_construct(
512
+ components={address_with_namespace(c, namespace) for c in self.components},
513
+ where_clause=(
514
+ self.where_clause.with_namespace(namespace)
515
+ if self.where_clause
516
+ else None
517
+ ),
518
+ )
519
+
520
+ @field_validator("components", mode="before")
521
+ def component_validator(cls, v, info: ValidationInfo):
522
+ output = set()
523
+ if isinstance(v, list):
524
+ for vc in v:
525
+ if isinstance(vc, Addressable):
526
+ output.add(vc._address)
527
+ else:
528
+ output.add(vc)
529
+ else:
530
+ output = v
531
+ if not isinstance(output, set):
532
+ raise ValueError(f"Invalid grain component {output}, is not set")
533
+ if not all(isinstance(x, str) for x in output):
534
+ raise ValueError(f"Invalid component {output}")
535
+ return output
536
+
537
+ def __add__(self, other: "Grain") -> "Grain":
538
+ if not other:
539
+ return self
540
+ where = self.where_clause
541
+ if other.where_clause:
542
+ if not self.where_clause:
543
+ where = other.where_clause
544
+ elif not other.where_clause == self.where_clause:
545
+ where = WhereClause.model_construct(
546
+ conditional=Conditional(
547
+ left=self.where_clause.conditional,
548
+ right=other.where_clause.conditional,
549
+ operator=BooleanOperator.AND,
550
+ )
551
+ )
552
+ # raise NotImplementedError(
553
+ # f"Cannot merge grains with where clauses, self {self.where_clause} other {other.where_clause}"
554
+ # )
555
+ return Grain(
556
+ components=self.components.union(other.components), where_clause=where
557
+ )
558
+
559
+ def __sub__(self, other: "Grain") -> "Grain":
560
+ return Grain.model_construct(
561
+ components=self.components.difference(other.components),
562
+ where_clause=self.where_clause,
563
+ )
564
+
565
+ def _gen_abstract(self) -> bool:
566
+ return not self.components or all(
567
+ [c.endswith(ALL_ROWS_CONCEPT) for c in self.components]
568
+ )
569
+
570
+ @property
571
+ def abstract(self):
572
+ if not self._abstract:
573
+ self._abstract = self._gen_abstract()
574
+ return self._abstract
575
+
576
+ def __eq__(self, other: object):
577
+ if isinstance(other, list):
578
+ if all([isinstance(c, Concept) for c in other]):
579
+ return self.components == set([c.address for c in other])
580
+ return False
581
+ if not isinstance(other, Grain):
582
+ return False
583
+ if self.components == other.components:
584
+ return True
585
+ return False
586
+
587
+ def issubset(self, other: "Grain"):
588
+ return self.components.issubset(other.components)
589
+
590
+ def union(self, other: "Grain"):
591
+ addresses = self.components.union(other.components)
592
+ return Grain(components=addresses, where_clause=self.where_clause)
593
+
594
+ def isdisjoint(self, other: "Grain"):
595
+ return self.components.isdisjoint(other.components)
596
+
597
+ def intersection(self, other: "Grain") -> "Grain":
598
+ intersection = self.components.intersection(other.components)
599
+ return Grain(components=intersection)
600
+
601
+ def _gen_str(self) -> str:
602
+ if self.abstract:
603
+ base = "Grain<Abstract>"
604
+ else:
605
+ base = "Grain<" + ",".join(sorted(self.components)) + ">"
606
+ if self.where_clause:
607
+ base += f"|{str(self.where_clause)}"
608
+ return base
609
+
610
+ def __str__(self):
611
+ if not self._str:
612
+ self._str = self._gen_str()
613
+ return self._str
614
+
615
+ def __radd__(self, other) -> "Grain":
616
+ if other == 0:
617
+ return self
618
+ else:
619
+ return self.__add__(other)
620
+
621
+
622
+ class Comparison(ConceptArgs, Mergeable, DataTyped, Namespaced, BaseModel):
623
+ left: Union[
624
+ int,
625
+ str,
626
+ float,
627
+ list,
628
+ bool,
629
+ datetime,
630
+ date,
631
+ Function,
632
+ ConceptRef,
633
+ Conditional,
634
+ DataType,
635
+ Comparison,
636
+ FunctionCallWrapper,
637
+ Parenthetical,
638
+ MagicConstants,
639
+ WindowItem,
640
+ AggregateWrapper,
641
+ FilterItem,
642
+ ]
643
+ right: Union[
644
+ int,
645
+ str,
646
+ float,
647
+ list,
648
+ bool,
649
+ date,
650
+ datetime,
651
+ ConceptRef,
652
+ Function,
653
+ Conditional,
654
+ DataType,
655
+ Comparison,
656
+ FunctionCallWrapper,
657
+ Parenthetical,
658
+ MagicConstants,
659
+ WindowItem,
660
+ AggregateWrapper,
661
+ TupleWrapper,
662
+ FilterItem,
663
+ ]
664
+ operator: ComparisonOperator
665
+
666
+ @field_validator("left", mode="before")
667
+ @classmethod
668
+ def left_validator(cls, v, info: ValidationInfo):
669
+ if isinstance(v, Concept):
670
+ return v.reference
671
+ return v
672
+
673
+ @field_validator("right", mode="before")
674
+ @classmethod
675
+ def right_validator(cls, v, info: ValidationInfo):
676
+ if isinstance(v, Concept):
677
+ return v.reference
678
+ return v
679
+
680
+ @model_validator(mode="after")
681
+ def validate_comparison(self):
682
+ left_type = arg_to_datatype(self.left)
683
+ right_type = arg_to_datatype(self.right)
684
+ left_name = (
685
+ left_type.name if isinstance(left_type, DataType) else str(left_type)
686
+ )
687
+ right_name = (
688
+ right_type.name if isinstance(right_type, DataType) else str(right_type)
689
+ )
690
+ if self.operator in (ComparisonOperator.IS, ComparisonOperator.IS_NOT):
691
+ if self.right != MagicConstants.NULL and DataType.BOOL != right_type:
692
+ raise SyntaxError(
693
+ f"Cannot use {self.operator.value} with non-null or boolean value {self.right}"
694
+ )
695
+ elif self.operator in (ComparisonOperator.IN, ComparisonOperator.NOT_IN):
696
+
697
+ if isinstance(right_type, ArrayType) and not is_compatible_datatype(
698
+ left_type, right_type.value_data_type
699
+ ):
700
+ raise SyntaxError(
701
+ f"Cannot compare {left_type} and {right_type} with operator {self.operator} in {str(self)}"
702
+ )
703
+ elif isinstance(self.right, Concept) and not is_compatible_datatype(
704
+ left_type, right_type
705
+ ):
706
+ raise SyntaxError(
707
+ f"Cannot compare {left_name} and {right_name} with operator {self.operator} in {str(self)}"
708
+ )
709
+ else:
710
+ if not is_compatible_datatype(left_type, right_type):
711
+ raise SyntaxError(
712
+ f"Cannot compare {left_name} ({self.left}) and {right_name} ({self.right}) of different types with operator {self.operator.value} in {str(self)}"
713
+ )
714
+
715
+ return self
716
+
717
+ def __add__(self, other):
718
+ if other is None:
719
+ return self
720
+ if not isinstance(other, (Comparison, Conditional, Parenthetical)):
721
+ raise ValueError("Cannot add Comparison to non-Comparison")
722
+ if other == self:
723
+ return self
724
+ return Conditional(left=self, right=other, operator=BooleanOperator.AND)
725
+
726
+ def __repr__(self):
727
+ if isinstance(self.left, Concept):
728
+ left = self.left.address
729
+ else:
730
+ left = str(self.left)
731
+ if isinstance(self.right, Concept):
732
+ right = self.right.address
733
+ else:
734
+ right = str(self.right)
735
+ return f"{left} {self.operator.value} {right}"
736
+
737
+ def __str__(self):
738
+ return self.__repr__()
739
+
740
+ def __eq__(self, other):
741
+ if not isinstance(other, Comparison):
742
+ return False
743
+ return (
744
+ self.left == other.left
745
+ and self.right == other.right
746
+ and self.operator == other.operator
747
+ )
748
+
749
+ def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
750
+ return self.__class__.model_construct(
751
+ left=(
752
+ self.left.with_merge(source, target, modifiers)
753
+ if isinstance(self.left, Mergeable)
754
+ else self.left
755
+ ),
756
+ right=(
757
+ self.right.with_merge(source, target, modifiers)
758
+ if isinstance(self.right, Mergeable)
759
+ else self.right
760
+ ),
761
+ operator=self.operator,
762
+ )
763
+
764
+ def with_reference_replacement(self, source, target):
765
+ return self.__class__.model_construct(
766
+ left=(
767
+ self.left.with_reference_replacement(source, target)
768
+ if isinstance(self.left, Mergeable)
769
+ else self.left
770
+ ),
771
+ right=(
772
+ self.right.with_reference_replacement(source, target)
773
+ if isinstance(self.right, Mergeable)
774
+ else self.right
775
+ ),
776
+ operator=self.operator,
777
+ )
778
+
779
+ def with_namespace(self, namespace: str):
780
+ return self.__class__.model_construct(
781
+ left=(
782
+ self.left.with_namespace(namespace)
783
+ if isinstance(self.left, Namespaced)
784
+ else self.left
785
+ ),
786
+ right=(
787
+ self.right.with_namespace(namespace)
788
+ if isinstance(self.right, Namespaced)
789
+ else self.right
790
+ ),
791
+ operator=self.operator,
792
+ )
793
+
794
+ @property
795
+ def concept_arguments(self) -> List[ConceptRef]:
796
+ """Return concepts directly referenced in where clause"""
797
+ output = []
798
+ output += get_concept_arguments(self.left)
799
+ output += get_concept_arguments(self.right)
800
+ return output
801
+
802
+ @property
803
+ def row_arguments(self) -> List[ConceptRef]:
804
+ output = []
805
+ output += get_concept_row_arguments(self.left)
806
+ output += get_concept_row_arguments(self.right)
807
+ return output
808
+
809
+ @property
810
+ def existence_arguments(self) -> List[Tuple[ConceptRef, ...]]:
811
+ """Return concepts directly referenced in where clause"""
812
+ output: List[Tuple[ConceptRef, ...]] = []
813
+ if isinstance(self.left, ConceptArgs):
814
+ output += self.left.existence_arguments
815
+ if isinstance(self.right, ConceptArgs):
816
+ output += self.right.existence_arguments
817
+ return output
818
+
819
+ @property
820
+ def output_datatype(self):
821
+ # a conditional is always a boolean
822
+ return DataType.BOOL
823
+
824
+
825
+ class SubselectComparison(Comparison):
826
+ def __eq__(self, other):
827
+ if not isinstance(other, SubselectComparison):
828
+ return False
829
+
830
+ comp = (
831
+ self.left == other.left
832
+ and self.right == other.right
833
+ and self.operator == other.operator
834
+ )
835
+ return comp
836
+
837
+ @property
838
+ def row_arguments(self) -> List[ConceptRef]:
839
+ return get_concept_row_arguments(self.left)
840
+
841
+ @property
842
+ def existence_arguments(self) -> list[tuple["ConceptRef", ...]]:
843
+ return [tuple(get_concept_arguments(self.right))]
844
+
845
+
846
+ class Concept(Addressable, DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
847
+ model_config = ConfigDict(
848
+ extra="forbid",
849
+ )
850
+ name: str
851
+ datatype: DataType | TraitDataType | ArrayType | StructType | MapType | NumericType
852
+ purpose: Purpose
853
+ derivation: Derivation = Derivation.ROOT
854
+ granularity: Granularity = Granularity.MULTI_ROW
855
+ metadata: Metadata = Field(
856
+ default_factory=lambda: Metadata(description=None, line_number=None),
857
+ validate_default=True,
858
+ )
859
+ lineage: Optional[
860
+ Union[
861
+ Function,
862
+ WindowItem,
863
+ FilterItem,
864
+ AggregateWrapper,
865
+ RowsetItem,
866
+ MultiSelectLineage,
867
+ Comparison,
868
+ ]
869
+ ] = None
870
+ namespace: str = Field(default=DEFAULT_NAMESPACE, validate_default=True)
871
+ keys: Optional[set[str]] = None
872
+ grain: "Grain" = Field(default=None, validate_default=True) # type: ignore
873
+ modifiers: List[Modifier] = Field(default_factory=list) # type: ignore
874
+ pseudonyms: set[str] = Field(default_factory=set)
875
+
876
+ def duplicate(self) -> Concept:
877
+ return self.model_copy(deep=True)
878
+
879
+ def __hash__(self):
880
+ return hash(
881
+ f"{self.name}+{self.datatype}+ {self.purpose} + {str(self.lineage)} + {self.namespace} + {str(self.grain)} + {str(self.keys)}"
882
+ )
883
+
884
+ def __repr__(self):
885
+ base = f"{self.address}@{self.grain}"
886
+ return base
887
+
888
+ @property
889
+ def is_internal(self) -> bool:
890
+ return self.namespace.startswith("_") or self.name.startswith("_")
891
+
892
+ @property
893
+ def reference(self) -> ConceptRef:
894
+ return ConceptRef.model_construct(
895
+ address=self.address,
896
+ datatype=self.output_datatype,
897
+ metadata=self.metadata,
898
+ )
899
+
900
+ @property
901
+ def output_datatype(self):
902
+ return self.datatype
903
+
904
+ @classmethod
905
+ def calculate_is_aggregate(cls, lineage):
906
+ if lineage and isinstance(lineage, Function):
907
+ if lineage.operator in FunctionClass.AGGREGATE_FUNCTIONS.value:
908
+ return True
909
+ if (
910
+ lineage
911
+ and isinstance(lineage, AggregateWrapper)
912
+ and lineage.function.operator in FunctionClass.AGGREGATE_FUNCTIONS.value
913
+ ):
914
+ return True
915
+ return False
916
+
917
+ @property
918
+ def is_aggregate(self):
919
+ base = getattr(self, "_is_aggregate", None)
920
+ if base:
921
+ return base
922
+ setattr(self, "_is_aggregate", self.calculate_is_aggregate(self.lineage))
923
+ return self._is_aggregate
924
+
925
+ def with_merge(self, source: Self, target: Self, modifiers: List[Modifier]) -> Self:
926
+ if self.address == source.address:
927
+ new = target.with_grain(self.grain.with_merge(source, target, modifiers))
928
+ new.pseudonyms.add(self.address)
929
+ return new
930
+ if not self.grain.components and not self.lineage and not self.keys:
931
+ return self
932
+ return self.__class__.model_construct(
933
+ name=self.name,
934
+ datatype=self.datatype,
935
+ purpose=self.purpose,
936
+ metadata=self.metadata,
937
+ derivation=self.derivation,
938
+ granularity=self.granularity,
939
+ lineage=(
940
+ self.lineage.with_merge(source, target, modifiers)
941
+ if self.lineage
942
+ else None
943
+ ),
944
+ grain=self.grain.with_merge(source, target, modifiers),
945
+ namespace=self.namespace,
946
+ keys=(
947
+ set(x if x != source.address else target.address for x in self.keys)
948
+ if self.keys
949
+ else None
950
+ ),
951
+ modifiers=self.modifiers,
952
+ pseudonyms=self.pseudonyms,
953
+ )
954
+
955
+ @field_validator("namespace", mode="plain")
956
+ @classmethod
957
+ def namespace_validation(cls, v):
958
+ return v or DEFAULT_NAMESPACE
959
+
960
+ @field_validator("metadata", mode="before")
961
+ @classmethod
962
+ def metadata_validation(cls, v):
963
+ v = v or Metadata()
964
+ return v
965
+
966
+ @field_validator("purpose", mode="after")
967
+ @classmethod
968
+ def purpose_validation(cls, v):
969
+ if v == Purpose.AUTO:
970
+ raise ValueError("Cannot set purpose to AUTO")
971
+ return v
972
+
973
+ @field_validator("grain", mode="before")
974
+ @classmethod
975
+ def parse_grain(cls, v, info: ValidationInfo) -> Grain:
976
+
977
+ # this is silly - rethink how we do grains
978
+ values = info.data
979
+
980
+ if not v and values.get("purpose", None) == Purpose.KEY:
981
+ v = Grain(
982
+ components={
983
+ f'{values.get("namespace", DEFAULT_NAMESPACE)}.{values["name"]}'
984
+ }
985
+ )
986
+ elif not v and values.get("purpose", None) == Purpose.PROPERTY:
987
+ v = Grain(components=values.get("keys", set()) or set())
988
+ elif (
989
+ "lineage" in values
990
+ and isinstance(values["lineage"], AggregateWrapper)
991
+ and values["lineage"].by
992
+ ):
993
+ v = Grain(components={c.address for c in values["lineage"].by})
994
+ elif not v:
995
+ v = Grain(components=set())
996
+ elif isinstance(v, Grain):
997
+ pass
998
+ elif isinstance(v, Concept):
999
+ v = Grain(components={v.address})
1000
+ elif isinstance(v, dict):
1001
+ v = Grain.model_validate(v)
1002
+ else:
1003
+ raise SyntaxError(f"Invalid grain {v} for concept {values['name']}")
1004
+ return v
1005
+
1006
+ def __eq__(self, other: object):
1007
+ if isinstance(other, str):
1008
+ if self.address == other:
1009
+ return True
1010
+ if isinstance(other, ConceptRef):
1011
+ return self.address == other.address
1012
+ if not isinstance(other, Concept):
1013
+ return False
1014
+ return (
1015
+ self.name == other.name
1016
+ and self.datatype == other.datatype
1017
+ and self.purpose == other.purpose
1018
+ and self.namespace == other.namespace
1019
+ and self.grain == other.grain
1020
+ and self.derivation == other.derivation
1021
+ and self.granularity == other.granularity
1022
+ # and self.keys == other.keys
1023
+ )
1024
+
1025
+ def __str__(self):
1026
+ grain = str(self.grain) if self.grain else "Grain<>"
1027
+ return f"{self.namespace}.{self.name}@{grain}"
1028
+
1029
+ @property
1030
+ def address(self) -> str:
1031
+ return f"{self.namespace}.{self.name}"
1032
+
1033
+ @property
1034
+ def output(self) -> "Concept":
1035
+ return self
1036
+
1037
+ @property
1038
+ def safe_address(self) -> str:
1039
+ if self.namespace == DEFAULT_NAMESPACE:
1040
+ return self.name.replace(".", "_")
1041
+ elif self.namespace:
1042
+ return f"{self.namespace.replace('.','_')}_{self.name.replace('.','_')}"
1043
+ return self.name.replace(".", "_")
1044
+
1045
+ def with_namespace(self, namespace: str) -> Self:
1046
+ return self.__class__.model_construct(
1047
+ name=self.name,
1048
+ datatype=self.datatype,
1049
+ purpose=self.purpose,
1050
+ granularity=self.granularity,
1051
+ derivation=self.derivation,
1052
+ metadata=self.metadata,
1053
+ lineage=self.lineage.with_namespace(namespace) if self.lineage else None,
1054
+ grain=(
1055
+ self.grain.with_namespace(namespace)
1056
+ if self.grain
1057
+ else Grain(components=set())
1058
+ ),
1059
+ namespace=(
1060
+ namespace + "." + self.namespace
1061
+ if self.namespace != DEFAULT_NAMESPACE
1062
+ else namespace
1063
+ ),
1064
+ keys=(
1065
+ set([address_with_namespace(x, namespace) for x in self.keys])
1066
+ if self.keys
1067
+ else None
1068
+ ),
1069
+ modifiers=self.modifiers,
1070
+ pseudonyms={address_with_namespace(v, namespace) for v in self.pseudonyms},
1071
+ )
1072
+
1073
+ def get_select_grain_and_keys(
1074
+ self, grain: Grain, environment: Environment
1075
+ ) -> Tuple[
1076
+ Function
1077
+ | WindowItem
1078
+ | FilterItem
1079
+ | AggregateWrapper
1080
+ | RowsetItem
1081
+ | MultiSelectLineage
1082
+ | Comparison
1083
+ | None,
1084
+ Grain,
1085
+ set[str] | None,
1086
+ ]:
1087
+ new_lineage = self.lineage
1088
+ final_grain = grain if not self.grain.components else self.grain
1089
+ keys = self.keys
1090
+
1091
+ if self.is_aggregate and grain.components and isinstance(new_lineage, Function):
1092
+ grain_components: list[ConceptRef | Concept] = [
1093
+ environment.concepts[c].reference for c in grain.components
1094
+ ]
1095
+ new_lineage = AggregateWrapper.model_construct(
1096
+ function=new_lineage, by=grain_components
1097
+ )
1098
+ final_grain = grain
1099
+ keys = set(grain.components)
1100
+ elif (
1101
+ grain
1102
+ and new_lineage
1103
+ and isinstance(new_lineage, AggregateWrapper)
1104
+ and not new_lineage.by
1105
+ ):
1106
+ grain_components = [
1107
+ environment.concepts[c].reference for c in grain.components
1108
+ ]
1109
+ new_lineage = AggregateWrapper.model_construct(
1110
+ function=new_lineage.function, by=grain_components
1111
+ )
1112
+ final_grain = grain
1113
+ keys = set([x.address for x in new_lineage.by])
1114
+ elif self.derivation == Derivation.BASIC:
1115
+
1116
+ pkeys: set[str] = set()
1117
+ assert new_lineage
1118
+ for x_ref in new_lineage.concept_arguments:
1119
+ x = environment.concepts[x_ref.address]
1120
+ if isinstance(x, UndefinedConcept):
1121
+ continue
1122
+ _, _, parent_keys = x.get_select_grain_and_keys(grain, environment)
1123
+ if parent_keys:
1124
+ pkeys.update(parent_keys)
1125
+ raw_keys = pkeys
1126
+ # deduplicate
1127
+
1128
+ final_grain = Grain.from_concepts(raw_keys, environment)
1129
+ keys = final_grain.components
1130
+ return new_lineage, final_grain, keys
1131
+
1132
+ def set_select_grain(self, grain: Grain, environment: Environment) -> Self:
1133
+ """Assign a mutable concept the appropriate grain/keys for a select"""
1134
+ new_lineage, final_grain, keys = self.get_select_grain_and_keys(
1135
+ grain, environment
1136
+ )
1137
+ return self.__class__.model_construct(
1138
+ name=self.name,
1139
+ datatype=self.datatype,
1140
+ purpose=self.purpose,
1141
+ granularity=self.granularity,
1142
+ derivation=self.derivation,
1143
+ metadata=self.metadata,
1144
+ lineage=new_lineage,
1145
+ grain=final_grain,
1146
+ namespace=self.namespace,
1147
+ keys=keys,
1148
+ modifiers=self.modifiers,
1149
+ pseudonyms=self.pseudonyms,
1150
+ )
1151
+
1152
+ def with_grain(self, grain: Optional["Grain"] = None) -> Self:
1153
+
1154
+ return self.__class__.model_construct(
1155
+ name=self.name,
1156
+ datatype=self.datatype,
1157
+ purpose=self.purpose,
1158
+ metadata=self.metadata,
1159
+ granularity=self.granularity,
1160
+ derivation=self.derivation,
1161
+ lineage=self.lineage,
1162
+ grain=grain if grain else Grain.model_construct(components=set()),
1163
+ namespace=self.namespace,
1164
+ keys=self.keys,
1165
+ modifiers=self.modifiers,
1166
+ pseudonyms=self.pseudonyms,
1167
+ )
1168
+
1169
+ @cached_property
1170
+ def sources(self) -> List["ConceptRef"]:
1171
+ if self.lineage:
1172
+ output: List[ConceptRef] = []
1173
+
1174
+ def get_sources(
1175
+ expr: Union[
1176
+ Function,
1177
+ WindowItem,
1178
+ FilterItem,
1179
+ AggregateWrapper,
1180
+ RowsetItem,
1181
+ MultiSelectLineage,
1182
+ Comparison,
1183
+ ],
1184
+ output: List[ConceptRef],
1185
+ ):
1186
+
1187
+ for item in expr.concept_arguments:
1188
+ if isinstance(item, (ConceptRef,)):
1189
+ if item.address == self.address:
1190
+ raise SyntaxError(
1191
+ f"Concept {self.address} references itself"
1192
+ )
1193
+ output.append(item)
1194
+
1195
+ # output += item.sources
1196
+
1197
+ get_sources(self.lineage, output)
1198
+ return output
1199
+ return []
1200
+
1201
+ @property
1202
+ def concept_arguments(self) -> List[ConceptRef]:
1203
+ return self.lineage.concept_arguments if self.lineage else []
1204
+
1205
+ @classmethod
1206
+ def calculate_derivation(self, lineage, purpose: Purpose) -> Derivation:
1207
+ from trilogy.core.models.build import (
1208
+ BuildAggregateWrapper,
1209
+ BuildComparison,
1210
+ BuildFilterItem,
1211
+ BuildFunction,
1212
+ BuildMultiSelectLineage,
1213
+ BuildRowsetItem,
1214
+ BuildWindowItem,
1215
+ )
1216
+
1217
+ if lineage and isinstance(lineage, (BuildWindowItem, WindowItem)):
1218
+ return Derivation.WINDOW
1219
+ elif lineage and isinstance(lineage, (BuildFilterItem, FilterItem)):
1220
+ return Derivation.FILTER
1221
+ elif lineage and isinstance(lineage, (BuildAggregateWrapper, AggregateWrapper)):
1222
+ return Derivation.AGGREGATE
1223
+ # elif lineage and isinstance(lineage, (BuildParenthetical, Parenthetical)):
1224
+ # return Derivation.PARENTHETICAL
1225
+ elif lineage and isinstance(lineage, (BuildRowsetItem, RowsetItem)):
1226
+ return Derivation.ROWSET
1227
+ elif lineage and isinstance(lineage, BuildComparison):
1228
+ return Derivation.BASIC
1229
+ elif lineage and isinstance(
1230
+ lineage, (BuildMultiSelectLineage, MultiSelectLineage)
1231
+ ):
1232
+ return Derivation.MULTISELECT
1233
+ elif (
1234
+ lineage
1235
+ and isinstance(lineage, (BuildFunction, Function))
1236
+ and lineage.operator in FunctionClass.AGGREGATE_FUNCTIONS.value
1237
+ ):
1238
+ return Derivation.AGGREGATE
1239
+ elif (
1240
+ lineage
1241
+ and isinstance(lineage, (BuildFunction, Function))
1242
+ and lineage.operator in FunctionClass.ONE_TO_MANY.value
1243
+ ):
1244
+ return Derivation.UNNEST
1245
+ elif (
1246
+ lineage
1247
+ and isinstance(lineage, (BuildFunction, Function))
1248
+ and lineage.operator == FunctionType.RECURSE_EDGE
1249
+ ):
1250
+ return Derivation.RECURSIVE
1251
+ elif (
1252
+ lineage
1253
+ and isinstance(lineage, (BuildFunction, Function))
1254
+ and lineage.operator == FunctionType.UNION
1255
+ ):
1256
+ return Derivation.UNION
1257
+ elif (
1258
+ lineage
1259
+ and isinstance(lineage, (BuildFunction, Function))
1260
+ and lineage.operator == FunctionType.GROUP
1261
+ ):
1262
+ return Derivation.GROUP_TO
1263
+ elif (
1264
+ lineage
1265
+ and isinstance(lineage, (BuildFunction, Function))
1266
+ and lineage.operator == FunctionType.ALIAS
1267
+ ):
1268
+ return Derivation.BASIC
1269
+ elif (
1270
+ lineage
1271
+ and isinstance(lineage, (BuildFunction, Function))
1272
+ and lineage.operator in FunctionClass.SINGLE_ROW.value
1273
+ ):
1274
+ return Derivation.CONSTANT
1275
+
1276
+ elif lineage and isinstance(lineage, (BuildFunction, Function)):
1277
+ if not lineage.concept_arguments:
1278
+ return Derivation.CONSTANT
1279
+ elif all(
1280
+ [x.derivation == Derivation.CONSTANT for x in lineage.concept_arguments]
1281
+ ):
1282
+ return Derivation.CONSTANT
1283
+ return Derivation.BASIC
1284
+ elif purpose == Purpose.CONSTANT:
1285
+ return Derivation.CONSTANT
1286
+ return Derivation.ROOT
1287
+
1288
+ @classmethod
1289
+ def calculate_granularity(cls, derivation: Derivation, grain: Grain, lineage):
1290
+ from trilogy.core.models.build import BuildFunction
1291
+
1292
+ if derivation == Derivation.CONSTANT:
1293
+ return Granularity.SINGLE_ROW
1294
+ elif derivation == Derivation.AGGREGATE:
1295
+ if all([x.endswith(ALL_ROWS_CONCEPT) for x in grain.components]):
1296
+ return Granularity.SINGLE_ROW
1297
+ elif (
1298
+ lineage
1299
+ and isinstance(lineage, (Function, BuildFunction))
1300
+ and lineage.operator
1301
+ in (FunctionType.UNNEST, FunctionType.UNION, FunctionType.DATE_SPINE)
1302
+ ):
1303
+ return Granularity.MULTI_ROW
1304
+ elif lineage and all(
1305
+ [x.granularity == Granularity.SINGLE_ROW for x in lineage.concept_arguments]
1306
+ ):
1307
+ return Granularity.SINGLE_ROW
1308
+ return Granularity.MULTI_ROW
1309
+
1310
+ # @property
1311
+ # def granularity(self) -> Granularity:
1312
+ # return self.calculate_granularity(self.derivation, self.grain, self.lineage)
1313
+
1314
+ def with_filter(
1315
+ self,
1316
+ condition: Conditional | Comparison | Parenthetical,
1317
+ environment: Environment | None = None,
1318
+ ) -> "Concept":
1319
+ from trilogy.utility import string_to_hash
1320
+
1321
+ if self.lineage and isinstance(self.lineage, FilterItem):
1322
+ if self.lineage.where.conditional == condition:
1323
+ return self
1324
+ hash = string_to_hash(self.name + str(condition))
1325
+ new_lineage = FilterItem(
1326
+ content=self.reference, where=WhereClause(conditional=condition)
1327
+ )
1328
+ new = Concept.model_construct(
1329
+ name=f"{self.name}_filter_{hash}",
1330
+ datatype=self.datatype,
1331
+ purpose=self.purpose,
1332
+ derivation=self.calculate_derivation(new_lineage, self.purpose),
1333
+ granularity=self.granularity,
1334
+ metadata=self.metadata,
1335
+ lineage=new_lineage,
1336
+ keys=(self.keys if self.purpose == Purpose.PROPERTY else None),
1337
+ grain=self.grain if self.grain else Grain(components=set()),
1338
+ namespace=self.namespace,
1339
+ modifiers=self.modifiers,
1340
+ pseudonyms=self.pseudonyms,
1341
+ )
1342
+ if environment:
1343
+ environment.add_concept(new)
1344
+ return new
1345
+
1346
+
1347
+ class UndefinedConceptFull(Concept, Mergeable, Namespaced):
1348
+ model_config = ConfigDict(arbitrary_types_allowed=True)
1349
+ name: str
1350
+ line_no: int | None = None
1351
+ datatype: (
1352
+ DataType | TraitDataType | ArrayType | StructType | MapType | NumericType
1353
+ ) = DataType.UNKNOWN
1354
+ purpose: Purpose = Purpose.UNKNOWN
1355
+
1356
+ @property
1357
+ def reference(self) -> UndefinedConcept:
1358
+ return UndefinedConcept(address=self.address)
1359
+
1360
+
1361
+ class OrderItem(Mergeable, ConceptArgs, Namespaced, BaseModel):
1362
+ # this needs to be a full concept as it may not exist in environment
1363
+ expr: Expr
1364
+ order: Ordering
1365
+
1366
+ @field_validator("expr", mode="before")
1367
+ def enforce_reference(cls, v):
1368
+ if isinstance(v, Concept):
1369
+ return v.reference
1370
+ return v
1371
+
1372
+ def with_namespace(self, namespace: str) -> "OrderItem":
1373
+ return OrderItem.model_construct(
1374
+ expr=(
1375
+ self.expr.with_namespace(namespace)
1376
+ if isinstance(self.expr, Namespaced)
1377
+ else self.expr
1378
+ ),
1379
+ order=self.order,
1380
+ )
1381
+
1382
+ def with_merge(
1383
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
1384
+ ) -> "OrderItem":
1385
+ return OrderItem.model_construct(
1386
+ expr=(
1387
+ self.expr.with_merge(source, target, modifiers)
1388
+ if isinstance(self.expr, Mergeable)
1389
+ else self.expr
1390
+ ),
1391
+ order=self.order,
1392
+ )
1393
+
1394
+ def with_reference_replacement(self, source, target):
1395
+ return OrderItem.model_construct(
1396
+ expr=(
1397
+ self.expr.with_reference_replacement(source, target)
1398
+ if isinstance(self.expr, Mergeable)
1399
+ else self.expr
1400
+ ),
1401
+ order=self.order,
1402
+ )
1403
+
1404
+ @property
1405
+ def concept_arguments(self) -> Sequence[ConceptRef]:
1406
+ return get_concept_arguments(self.expr)
1407
+
1408
+ @property
1409
+ def row_arguments(self) -> Sequence[ConceptRef]:
1410
+ if isinstance(self.expr, ConceptArgs):
1411
+ return self.expr.row_arguments
1412
+ return self.concept_arguments
1413
+
1414
+ @property
1415
+ def existence_arguments(self) -> Sequence[tuple["ConceptRef", ...]]:
1416
+ if isinstance(self.expr, ConceptArgs):
1417
+ return self.expr.existence_arguments
1418
+ return []
1419
+
1420
+ @property
1421
+ def output_datatype(self):
1422
+ return arg_to_datatype(self.expr)
1423
+
1424
+
1425
+ class WindowItem(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
1426
+ type: WindowType
1427
+ content: FuncArgs
1428
+ order_by: List["OrderItem"]
1429
+ over: List["ConceptRef"] = Field(default_factory=list)
1430
+ index: Optional[int] = None
1431
+
1432
+ def __str__(self):
1433
+ return self.__repr__()
1434
+
1435
+ def __repr__(self):
1436
+ return f"{self.type.value} {self.content} by {self.index} over {self.over} order {self.order_by}"
1437
+
1438
+ @field_validator("content", mode="before")
1439
+ def enforce_concept_ref(cls, v):
1440
+ if isinstance(v, Concept):
1441
+ return ConceptRef(address=v.address, datatype=v.datatype)
1442
+ return v
1443
+
1444
+ @field_validator("over", mode="before")
1445
+ def enforce_concept_ref_over(cls, v):
1446
+ final = []
1447
+ for item in v:
1448
+ if isinstance(item, Concept):
1449
+ final.append(ConceptRef(address=item.address, datatype=item.datatype))
1450
+ else:
1451
+ final.append(item)
1452
+ return final
1453
+
1454
+ def with_merge(
1455
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
1456
+ ) -> "WindowItem":
1457
+ output = WindowItem.model_construct(
1458
+ type=self.type,
1459
+ content=(
1460
+ self.content.with_merge(source, target, modifiers)
1461
+ if isinstance(self.content, Mergeable)
1462
+ else self.content
1463
+ ),
1464
+ over=[x.with_merge(source, target, modifiers) for x in self.over],
1465
+ order_by=[x.with_merge(source, target, modifiers) for x in self.order_by],
1466
+ index=self.index,
1467
+ )
1468
+ return output
1469
+
1470
+ def with_reference_replacement(self, source, target):
1471
+ return WindowItem.model_construct(
1472
+ type=self.type,
1473
+ content=self.content.with_reference_replacement(source, target),
1474
+ over=[x.with_reference_replacement(source, target) for x in self.over],
1475
+ order_by=[
1476
+ x.with_reference_replacement(source, target) for x in self.order_by
1477
+ ],
1478
+ index=self.index,
1479
+ )
1480
+
1481
+ def with_namespace(self, namespace: str) -> "WindowItem":
1482
+ return WindowItem.model_construct(
1483
+ type=self.type,
1484
+ content=(
1485
+ self.content.with_namespace(namespace)
1486
+ if isinstance(self.content, Namespaced)
1487
+ else self.content
1488
+ ),
1489
+ over=[x.with_namespace(namespace) for x in self.over],
1490
+ order_by=[x.with_namespace(namespace) for x in self.order_by],
1491
+ index=self.index,
1492
+ )
1493
+
1494
+ @property
1495
+ def concept_arguments(self) -> List[ConceptRef]:
1496
+ output = []
1497
+ output += get_concept_arguments(self.content)
1498
+ for order in self.order_by:
1499
+ output += get_concept_arguments(order)
1500
+ for item in self.over:
1501
+ output += get_concept_arguments(item)
1502
+ return output
1503
+
1504
+ @property
1505
+ def output_datatype(self):
1506
+ if self.type in (WindowType.RANK, WindowType.ROW_NUMBER):
1507
+ return DataType.INTEGER
1508
+ return self.content.output_datatype
1509
+
1510
+
1511
+ def get_basic_type(
1512
+ type: DataType | ArrayType | StructType | MapType | NumericType | TraitDataType,
1513
+ ) -> DataType:
1514
+ if isinstance(type, ArrayType):
1515
+ return DataType.ARRAY
1516
+ if isinstance(type, StructType):
1517
+ return DataType.STRUCT
1518
+ if isinstance(type, MapType):
1519
+ return DataType.MAP
1520
+ if isinstance(type, NumericType):
1521
+ return DataType.NUMERIC
1522
+ if isinstance(type, TraitDataType):
1523
+ return get_basic_type(type.type)
1524
+ return type
1525
+
1526
+
1527
+ class CaseWhen(Namespaced, DataTyped, ConceptArgs, Mergeable, BaseModel):
1528
+ comparison: Conditional | SubselectComparison | Comparison
1529
+ expr: "Expr"
1530
+
1531
+ @field_validator("expr", mode="before")
1532
+ def enforce_reference(cls, v):
1533
+ if isinstance(v, Concept):
1534
+ return v.reference
1535
+ return v
1536
+
1537
+ @property
1538
+ def output_datatype(self):
1539
+ return arg_to_datatype(self.expr)
1540
+
1541
+ def __str__(self):
1542
+ return self.__repr__()
1543
+
1544
+ def __repr__(self):
1545
+ return f"WHEN {str(self.comparison)} THEN {str(self.expr)}"
1546
+
1547
+ @property
1548
+ def concept_arguments(self):
1549
+ return get_concept_arguments(self.comparison) + get_concept_arguments(self.expr)
1550
+
1551
+ @property
1552
+ def concept_row_arguments(self):
1553
+ return get_concept_row_arguments(self.comparison) + get_concept_row_arguments(
1554
+ self.expr
1555
+ )
1556
+
1557
+ def with_namespace(self, namespace: str) -> CaseWhen:
1558
+ return CaseWhen.model_construct(
1559
+ comparison=self.comparison.with_namespace(namespace),
1560
+ expr=(
1561
+ self.expr.with_namespace(namespace)
1562
+ if isinstance(
1563
+ self.expr,
1564
+ Namespaced,
1565
+ )
1566
+ else self.expr
1567
+ ),
1568
+ )
1569
+
1570
+ def with_merge(
1571
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
1572
+ ) -> CaseWhen:
1573
+ return CaseWhen.model_construct(
1574
+ comparison=self.comparison.with_merge(source, target, modifiers),
1575
+ expr=(
1576
+ self.expr.with_merge(source, target, modifiers)
1577
+ if isinstance(self.expr, Mergeable)
1578
+ else self.expr
1579
+ ),
1580
+ )
1581
+
1582
+ def with_reference_replacement(self, source, target):
1583
+ return CaseWhen.model_construct(
1584
+ comparison=self.comparison.with_reference_replacement(source, target),
1585
+ expr=(
1586
+ self.expr.with_reference_replacement(source, target)
1587
+ if isinstance(self.expr, Mergeable)
1588
+ else self.expr
1589
+ ),
1590
+ )
1591
+
1592
+
1593
+ class CaseElse(Namespaced, ConceptArgs, DataTyped, Mergeable, BaseModel):
1594
+ expr: "Expr"
1595
+ # this ensures that it's easily differentiable from CaseWhen
1596
+ discriminant: ComparisonOperator = ComparisonOperator.ELSE
1597
+
1598
+ def __str__(self):
1599
+ return self.__repr__()
1600
+
1601
+ def __repr__(self):
1602
+ return f"ELSE {str(self.expr)}"
1603
+
1604
+ @property
1605
+ def output_datatype(self):
1606
+ return arg_to_datatype(self.expr)
1607
+
1608
+ @field_validator("expr", mode="before")
1609
+ def enforce_expr(cls, v):
1610
+ if isinstance(v, Concept):
1611
+ return ConceptRef(address=v.address, datatype=v.datatype)
1612
+ return v
1613
+
1614
+ @property
1615
+ def concept_arguments(self):
1616
+ return get_concept_arguments(self.expr)
1617
+
1618
+ def with_merge(
1619
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
1620
+ ) -> CaseElse:
1621
+ return CaseElse.model_construct(
1622
+ discriminant=self.discriminant,
1623
+ expr=(
1624
+ self.expr.with_merge(source, target, modifiers)
1625
+ if isinstance(self.expr, Mergeable)
1626
+ else self.expr
1627
+ ),
1628
+ )
1629
+
1630
+ def with_reference_replacement(self, source, target):
1631
+ return CaseElse.model_construct(
1632
+ discriminant=self.discriminant,
1633
+ expr=(
1634
+ self.expr.with_reference_replacement(
1635
+ source,
1636
+ target,
1637
+ )
1638
+ if isinstance(self.expr, Mergeable)
1639
+ else self.expr
1640
+ ),
1641
+ )
1642
+
1643
+ def with_namespace(self, namespace: str) -> CaseElse:
1644
+ return CaseElse.model_construct(
1645
+ discriminant=self.discriminant,
1646
+ expr=(
1647
+ self.expr.with_namespace(namespace)
1648
+ if isinstance(
1649
+ self.expr,
1650
+ Namespaced,
1651
+ )
1652
+ else self.expr
1653
+ ),
1654
+ )
1655
+
1656
+
1657
+ def get_concept_row_arguments(expr) -> List["ConceptRef"]:
1658
+ output = []
1659
+ if isinstance(expr, ConceptRef):
1660
+ output += [expr]
1661
+
1662
+ elif isinstance(expr, ConceptArgs):
1663
+ output += expr.row_arguments
1664
+ return output
1665
+
1666
+
1667
+ def get_concept_arguments(expr) -> List["ConceptRef"]:
1668
+ output = []
1669
+ if isinstance(expr, ConceptRef):
1670
+ output += [expr]
1671
+
1672
+ elif isinstance(
1673
+ expr,
1674
+ ConceptArgs,
1675
+ ):
1676
+ output += expr.concept_arguments
1677
+ return output
1678
+
1679
+
1680
+ def args_to_pretty(input: set[DataType]) -> str:
1681
+ return ", ".join(sorted([f"'{x.value}'" for x in input if x != DataType.UNKNOWN]))
1682
+
1683
+
1684
+ class Function(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
1685
+ operator: FunctionType
1686
+ arg_count: int = Field(default=1)
1687
+ output_datatype: (
1688
+ DataType | ArrayType | StructType | MapType | NumericType | TraitDataType
1689
+ )
1690
+ output_purpose: Purpose
1691
+ valid_inputs: Optional[
1692
+ Union[
1693
+ Set[DataType],
1694
+ List[Set[DataType]],
1695
+ ]
1696
+ ] = None
1697
+ arguments: Sequence[FuncArgs]
1698
+
1699
+ class Config:
1700
+ frozen = True
1701
+
1702
+ def __repr__(self):
1703
+ return f'{self.operator.value}({",".join([str(a) for a in self.arguments])})'
1704
+
1705
+ def __str__(self):
1706
+ return self.__repr__()
1707
+
1708
+ @property
1709
+ def datatype(self):
1710
+ return self.output_datatype
1711
+
1712
+ @field_validator("arguments", mode="before")
1713
+ @classmethod
1714
+ def parse_arguments(cls, v, info: ValidationInfo):
1715
+ from trilogy.core.models.build import BuildConcept
1716
+ from trilogy.parsing.exceptions import ParseError
1717
+
1718
+ values = info.data
1719
+ arg_count = len(v)
1720
+ final = []
1721
+ for x in v:
1722
+ if isinstance(x, Concept) and not isinstance(x, BuildConcept):
1723
+ final.append(x.reference)
1724
+ else:
1725
+ final.append(x)
1726
+ v = final
1727
+ target_arg_count = values["arg_count"]
1728
+ operator_name = values["operator"].name
1729
+ # surface right error
1730
+ if "valid_inputs" not in values:
1731
+ return v
1732
+ valid_inputs = values["valid_inputs"]
1733
+ if not arg_count <= target_arg_count:
1734
+ if target_arg_count != InfiniteFunctionArgs:
1735
+ raise ParseError(
1736
+ f"Incorrect argument count to {operator_name} function, expects"
1737
+ f" {target_arg_count}, got {arg_count}"
1738
+ )
1739
+ # if all arguments can be any of the set type
1740
+ # turn this into an array for validation
1741
+ if isinstance(valid_inputs, set):
1742
+ valid_inputs = [valid_inputs for _ in v]
1743
+ elif not valid_inputs:
1744
+ return v
1745
+ for idx, arg in enumerate(v):
1746
+ if (
1747
+ isinstance(arg, ConceptRef)
1748
+ and get_basic_type(arg.datatype.data_type) not in valid_inputs[idx]
1749
+ ):
1750
+ if arg.datatype != DataType.UNKNOWN:
1751
+
1752
+ raise TypeError(
1753
+ f"Invalid argument type '{arg.datatype.data_type.value}' passed into {operator_name} function in position {idx+1}"
1754
+ f" from concept: {arg.name}. Valid: {args_to_pretty(valid_inputs[idx])}."
1755
+ )
1756
+ if (
1757
+ isinstance(arg, Function)
1758
+ and get_basic_type(arg.output_datatype) not in valid_inputs[idx]
1759
+ ):
1760
+ if arg.output_datatype != DataType.UNKNOWN:
1761
+ raise TypeError(
1762
+ f"Invalid argument type {arg.output_datatype}' passed into"
1763
+ f" {operator_name} function from function {arg.operator.name} in position {idx+1}. Valid: {args_to_pretty(valid_inputs[idx])}"
1764
+ )
1765
+ # check constants
1766
+ comparisons: List[Tuple[Type, DataType]] = [
1767
+ (str, DataType.STRING),
1768
+ (int, DataType.INTEGER),
1769
+ (float, DataType.FLOAT),
1770
+ (bool, DataType.BOOL),
1771
+ (DatePart, DataType.DATE_PART),
1772
+ ]
1773
+ for ptype, dtype in comparisons:
1774
+ if (
1775
+ isinstance(arg, ptype)
1776
+ and get_basic_type(dtype) in valid_inputs[idx]
1777
+ ):
1778
+ # attempt to exit early to avoid checking all types
1779
+ break
1780
+ elif isinstance(arg, ptype):
1781
+ if isinstance(arg, str) and DataType.DATE_PART in valid_inputs[idx]:
1782
+ if arg not in [x.value for x in DatePart]:
1783
+ pass
1784
+ else:
1785
+ break
1786
+ raise TypeError(
1787
+ f'Invalid {dtype} constant passed into {operator_name} "{arg}", expecting one of {valid_inputs[idx]}'
1788
+ )
1789
+ return v
1790
+
1791
+ def with_reference_replacement(self, source: str, target: Expr | ArgBinding):
1792
+ from trilogy.core.functions import arg_to_datatype, merge_datatypes
1793
+
1794
+ nargs = [
1795
+ (
1796
+ c.with_reference_replacement(
1797
+ source,
1798
+ target,
1799
+ )
1800
+ if isinstance(
1801
+ c,
1802
+ Mergeable,
1803
+ )
1804
+ else c
1805
+ )
1806
+ for c in self.arguments
1807
+ ]
1808
+ if self.output_datatype == DataType.UNKNOWN:
1809
+ new_output = merge_datatypes([arg_to_datatype(x) for x in nargs])
1810
+
1811
+ if self.operator == FunctionType.ATTR_ACCESS:
1812
+ if isinstance(new_output, StructType):
1813
+ new_output = new_output.field_types[str(nargs[1])]
1814
+ else:
1815
+ new_output = self.output_datatype
1816
+ # this is not ideal - see hacky logic for datatypes above
1817
+ # we need to figure out how to patch properly
1818
+ # should use function factory, but does not have environment access
1819
+ # probably move all datatype resolution to build?
1820
+ return Function.model_construct(
1821
+ operator=self.operator,
1822
+ arguments=nargs,
1823
+ output_datatype=new_output,
1824
+ output_purpose=self.output_purpose,
1825
+ valid_inputs=self.valid_inputs,
1826
+ arg_count=self.arg_count,
1827
+ )
1828
+
1829
+ def with_namespace(self, namespace: str) -> "Function":
1830
+ return Function.model_construct(
1831
+ operator=self.operator,
1832
+ arguments=[
1833
+ (
1834
+ c.with_namespace(namespace)
1835
+ if isinstance(
1836
+ c,
1837
+ Namespaced,
1838
+ )
1839
+ else c
1840
+ )
1841
+ for c in self.arguments
1842
+ ],
1843
+ output_datatype=self.output_datatype,
1844
+ output_purpose=self.output_purpose,
1845
+ valid_inputs=self.valid_inputs,
1846
+ arg_count=self.arg_count,
1847
+ )
1848
+
1849
+ def with_merge(
1850
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
1851
+ ) -> "Function":
1852
+ return Function.model_construct(
1853
+ operator=self.operator,
1854
+ arguments=[
1855
+ (
1856
+ c.with_merge(source, target, modifiers)
1857
+ if isinstance(
1858
+ c,
1859
+ Mergeable,
1860
+ )
1861
+ else c
1862
+ )
1863
+ for c in self.arguments
1864
+ ],
1865
+ output_datatype=self.output_datatype,
1866
+ output_purpose=self.output_purpose,
1867
+ valid_inputs=self.valid_inputs,
1868
+ arg_count=self.arg_count,
1869
+ )
1870
+
1871
+ @property
1872
+ def concept_arguments(self) -> List[ConceptRef]:
1873
+ base = []
1874
+ for arg in self.arguments:
1875
+ base += get_concept_arguments(arg)
1876
+ return base
1877
+
1878
+
1879
+ class FunctionCallWrapper(
1880
+ DataTyped,
1881
+ ConceptArgs,
1882
+ Mergeable,
1883
+ Namespaced,
1884
+ BaseModel,
1885
+ ):
1886
+ content: Expr
1887
+ name: str
1888
+ args: List[Expr]
1889
+
1890
+ def __str__(self):
1891
+ return f'@{self.name}({",".join([str(x) for x in self.args])})'
1892
+
1893
+ def with_namespace(self, namespace) -> "FunctionCallWrapper":
1894
+ return FunctionCallWrapper.model_construct(
1895
+ content=(
1896
+ self.content.with_namespace(namespace)
1897
+ if isinstance(self.content, Namespaced)
1898
+ else self.content
1899
+ ),
1900
+ name=self.name,
1901
+ args=[
1902
+ x.with_namespace(namespace) if isinstance(x, Namespaced) else x
1903
+ for x in self.args
1904
+ ],
1905
+ )
1906
+
1907
+ def with_reference_replacement(self, source, target):
1908
+ raise NotImplementedError("Cannot reference replace")
1909
+ return self
1910
+
1911
+ def with_merge(
1912
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
1913
+ ) -> "FunctionCallWrapper":
1914
+ return FunctionCallWrapper.model_construct(
1915
+ content=(
1916
+ self.content.with_merge(source, target, modifiers)
1917
+ if isinstance(self.content, Mergeable)
1918
+ else self.content
1919
+ ),
1920
+ name=self.name,
1921
+ args=[
1922
+ (
1923
+ x.with_merge(source, target, modifiers)
1924
+ if isinstance(x, Mergeable)
1925
+ else x
1926
+ )
1927
+ for x in self.args
1928
+ ],
1929
+ )
1930
+
1931
+ @property
1932
+ def concept_arguments(self) -> Sequence[ConceptRef]:
1933
+ base: List[ConceptRef] = []
1934
+ x = self.content
1935
+ if isinstance(x, ConceptRef):
1936
+ base += [x]
1937
+ elif isinstance(x, ConceptArgs):
1938
+ base += x.concept_arguments
1939
+ return base
1940
+
1941
+ @property
1942
+ def output_datatype(self):
1943
+ return arg_to_datatype(self.content)
1944
+
1945
+
1946
+ class AggregateWrapper(Mergeable, DataTyped, ConceptArgs, Namespaced, BaseModel):
1947
+ function: Function
1948
+ by: List[ConceptRef | Concept] = Field(default_factory=list)
1949
+
1950
+ @field_validator("by", mode="before")
1951
+ @classmethod
1952
+ def enforce_concept_ref(cls, v):
1953
+ output = []
1954
+ for item in v:
1955
+ if isinstance(item, Concept):
1956
+ output.append(item.reference)
1957
+ else:
1958
+ output.append(item)
1959
+ return output
1960
+
1961
+ def __str__(self):
1962
+ grain_str = [str(c) for c in self.by] if self.by else "abstract"
1963
+ return f"{str(self.function)}<{grain_str}>"
1964
+
1965
+ @property
1966
+ def datatype(self):
1967
+ return self.function.datatype
1968
+
1969
+ @property
1970
+ def concept_arguments(self) -> List[ConceptRef]:
1971
+ return self.function.concept_arguments + [x.reference for x in self.by]
1972
+
1973
+ @property
1974
+ def output_datatype(self):
1975
+ return self.function.output_datatype
1976
+
1977
+ @property
1978
+ def output_purpose(self):
1979
+ return self.function.output_purpose
1980
+
1981
+ def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
1982
+ return AggregateWrapper.model_construct(
1983
+ function=self.function.with_merge(source, target, modifiers=modifiers),
1984
+ by=(
1985
+ [c.with_merge(source, target, modifiers) for c in self.by]
1986
+ if self.by
1987
+ else []
1988
+ ),
1989
+ )
1990
+
1991
+ def with_reference_replacement(self, source, target):
1992
+ return AggregateWrapper.model_construct(
1993
+ function=self.function.with_reference_replacement(source, target),
1994
+ by=(
1995
+ [c.with_reference_replacement(source, target) for c in self.by]
1996
+ if self.by
1997
+ else []
1998
+ ),
1999
+ )
2000
+
2001
+ def with_namespace(self, namespace: str) -> "AggregateWrapper":
2002
+ return AggregateWrapper.model_construct(
2003
+ function=self.function.with_namespace(namespace),
2004
+ by=[c.with_namespace(namespace) for c in self.by] if self.by else [],
2005
+ )
2006
+
2007
+
2008
+ class FilterItem(DataTyped, Namespaced, ConceptArgs, BaseModel):
2009
+ content: FuncArgs
2010
+ where: "WhereClause"
2011
+
2012
+ @field_validator("content", mode="before")
2013
+ def enforce_concept_ref(cls, v):
2014
+ if isinstance(v, Concept):
2015
+ return ConceptRef(address=v.address, datatype=v.datatype)
2016
+ return v
2017
+
2018
+ def __str__(self):
2019
+ return f"<Filter: {str(self.content)} where {str(self.where)}>"
2020
+
2021
+ def with_merge(
2022
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
2023
+ ) -> "FilterItem":
2024
+ return FilterItem.model_construct(
2025
+ content=(
2026
+ self.content.with_merge(source, target, modifiers)
2027
+ if isinstance(self.content, Mergeable)
2028
+ else self.content
2029
+ ),
2030
+ where=self.where.with_merge(source, target, modifiers),
2031
+ )
2032
+
2033
+ def with_namespace(self, namespace: str) -> "FilterItem":
2034
+ return FilterItem.model_construct(
2035
+ content=(
2036
+ self.content.with_namespace(namespace)
2037
+ if isinstance(self.content, Namespaced)
2038
+ else self.content
2039
+ ),
2040
+ where=self.where.with_namespace(namespace),
2041
+ )
2042
+
2043
+ @property
2044
+ def output_datatype(self):
2045
+ return arg_to_datatype(self.content)
2046
+
2047
+ @property
2048
+ def concept_arguments(self):
2049
+ if isinstance(self.content, ConceptRef):
2050
+ return [self.content] + self.where.concept_arguments
2051
+ elif isinstance(self.content, ConceptArgs):
2052
+ return self.content.concept_arguments + self.where.concept_arguments
2053
+ return self.where.concept_arguments
2054
+
2055
+
2056
+ class RowsetLineage(Namespaced, Mergeable, BaseModel):
2057
+ name: str
2058
+ derived_concepts: List[ConceptRef]
2059
+ select: SelectLineage | MultiSelectLineage
2060
+
2061
+ def with_namespace(self, namespace: str):
2062
+ return RowsetLineage.model_construct(
2063
+ name=self.name,
2064
+ derived_concepts=[
2065
+ x.with_namespace(namespace) for x in self.derived_concepts
2066
+ ],
2067
+ select=self.select.with_namespace(namespace),
2068
+ )
2069
+
2070
+ def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
2071
+ return RowsetLineage.model_construct(
2072
+ name=self.name,
2073
+ derived_concepts=[
2074
+ x.with_merge(source, target, modifiers) for x in self.derived_concepts
2075
+ ],
2076
+ select=self.select.with_merge(source, target, modifiers),
2077
+ )
2078
+
2079
+
2080
+ class RowsetItem(Mergeable, DataTyped, ConceptArgs, Namespaced, BaseModel):
2081
+ content: ConceptRef
2082
+ rowset: RowsetLineage
2083
+
2084
+ def __repr__(self):
2085
+ return f"<Rowset<{self.rowset.name}>: {str(self.content)}>"
2086
+
2087
+ def __str__(self):
2088
+ return self.__repr__()
2089
+
2090
+ def with_merge(self, source: Concept, target: Concept, modifiers: List[Modifier]):
2091
+ return RowsetItem.model_construct(
2092
+ content=self.content.with_merge(source, target, modifiers),
2093
+ rowset=self.rowset,
2094
+ )
2095
+
2096
+ def with_namespace(self, namespace: str) -> "RowsetItem":
2097
+ return RowsetItem.model_construct(
2098
+ content=self.content.with_namespace(namespace),
2099
+ rowset=self.rowset.with_namespace(namespace),
2100
+ )
2101
+
2102
+ @property
2103
+ def output(self) -> ConceptRef:
2104
+ return self.content
2105
+
2106
+ @property
2107
+ def output_datatype(self):
2108
+ return self.content.datatype
2109
+
2110
+ @property
2111
+ def concept_arguments(self):
2112
+ return [self.content]
2113
+
2114
+
2115
+ class OrderBy(Mergeable, Namespaced, BaseModel):
2116
+ items: List[OrderItem]
2117
+
2118
+ def with_namespace(self, namespace: str) -> "OrderBy":
2119
+ return OrderBy.model_construct(
2120
+ items=[x.with_namespace(namespace) for x in self.items]
2121
+ )
2122
+
2123
+ def with_merge(
2124
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
2125
+ ) -> "OrderBy":
2126
+ return OrderBy.model_construct(
2127
+ items=[x.with_merge(source, target, modifiers) for x in self.items]
2128
+ )
2129
+
2130
+ @property
2131
+ def concept_arguments(self):
2132
+ base = []
2133
+ for x in self.items:
2134
+ base += x.concept_arguments
2135
+ return base
2136
+
2137
+
2138
+ class AlignClause(Namespaced, BaseModel):
2139
+ items: List[AlignItem]
2140
+
2141
+ def with_namespace(self, namespace: str) -> "AlignClause":
2142
+ return AlignClause.model_construct(
2143
+ items=[x.with_namespace(namespace) for x in self.items]
2144
+ )
2145
+
2146
+
2147
+ class DeriveItem(Namespaced, DataTyped, ConceptArgs, Mergeable, BaseModel):
2148
+ expr: Expr
2149
+ name: str
2150
+ namespace: str
2151
+
2152
+ @property
2153
+ def derived_concept(self) -> str:
2154
+ return f"{self.namespace}.{self.name}"
2155
+ # return ConceptRef(
2156
+ # address=f"{self.namespace}.{self.name}",
2157
+ # datatype=arg_to_datatype(self.expr),
2158
+ # )
2159
+
2160
+ def with_namespace(self, namespace):
2161
+ return DeriveItem.model_construct(
2162
+ expr=(self.expr.with_namespace(namespace) if self.expr else None),
2163
+ name=self.name,
2164
+ namespace=namespace,
2165
+ )
2166
+
2167
+ def with_merge(
2168
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
2169
+ ) -> "DeriveItem":
2170
+ return DeriveItem.model_construct(
2171
+ expr=(
2172
+ self.expr.with_merge(source, target, modifiers)
2173
+ if isinstance(self.expr, Mergeable)
2174
+ else self.expr
2175
+ ),
2176
+ name=self.name,
2177
+ namespace=self.namespace,
2178
+ )
2179
+
2180
+ def with_reference_replacement(self, source, target):
2181
+ return DeriveItem.model_construct(
2182
+ expr=(
2183
+ self.expr.with_reference_replacement(source, target)
2184
+ if isinstance(self.expr, Mergeable)
2185
+ else self.expr
2186
+ ),
2187
+ name=self.name,
2188
+ namespace=self.namespace,
2189
+ )
2190
+
2191
+
2192
+ class DeriveClause(Mergeable, Namespaced, BaseModel):
2193
+ items: List[DeriveItem]
2194
+
2195
+ def with_namespace(self, namespace: str) -> "DeriveClause":
2196
+ return DeriveClause.model_construct(
2197
+ items=[
2198
+ x.with_namespace(namespace) if isinstance(x, Namespaced) else x
2199
+ for x in self.items
2200
+ ]
2201
+ )
2202
+
2203
+ def with_merge(
2204
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
2205
+ ) -> "DeriveClause":
2206
+ return DeriveClause.model_construct(
2207
+ items=[
2208
+ (
2209
+ x.with_merge(source, target, modifiers)
2210
+ if isinstance(x, Mergeable)
2211
+ else x
2212
+ )
2213
+ for x in self.items
2214
+ ]
2215
+ )
2216
+
2217
+ def with_reference_replacement(self, source, target):
2218
+ return DeriveClause.model_construct(
2219
+ items=[
2220
+ (
2221
+ x.with_reference_replacement(source, target)
2222
+ if isinstance(x, Mergeable)
2223
+ else x
2224
+ )
2225
+ for x in self.items
2226
+ ]
2227
+ )
2228
+
2229
+
2230
+ class SelectLineage(Mergeable, Namespaced, BaseModel):
2231
+ selection: List[ConceptRef]
2232
+ hidden_components: set[str]
2233
+ local_concepts: dict[str, Concept]
2234
+ order_by: Optional[OrderBy] = None
2235
+ limit: Optional[int] = None
2236
+ meta: Metadata = Field(default_factory=lambda: Metadata())
2237
+ grain: Grain = Field(default_factory=Grain)
2238
+ where_clause: Union["WhereClause", None] = Field(default=None)
2239
+ having_clause: Union["HavingClause", None] = Field(default=None)
2240
+
2241
+ @property
2242
+ def output_components(self) -> List[ConceptRef]:
2243
+ return self.selection
2244
+
2245
+ def with_merge(
2246
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
2247
+ ) -> SelectLineage:
2248
+ return SelectLineage.model_construct(
2249
+ selection=[x.with_merge(source, target, modifiers) for x in self.selection],
2250
+ hidden_components=self.hidden_components,
2251
+ local_concepts={
2252
+ x: y.with_merge(source, target, modifiers)
2253
+ for x, y in self.local_concepts.items()
2254
+ },
2255
+ order_by=(
2256
+ self.order_by.with_merge(source, target, modifiers)
2257
+ if self.order_by
2258
+ else None
2259
+ ),
2260
+ limit=self.limit,
2261
+ grain=self.grain.with_merge(source, target, modifiers),
2262
+ where_clause=(
2263
+ self.where_clause.with_merge(source, target, modifiers)
2264
+ if self.where_clause
2265
+ else None
2266
+ ),
2267
+ having_clause=(
2268
+ self.having_clause.with_merge(source, target, modifiers)
2269
+ if self.having_clause
2270
+ else None
2271
+ ),
2272
+ )
2273
+
2274
+ def with_namespace(self, namespace):
2275
+ return SelectLineage.model_construct(
2276
+ selection=[x.with_namespace(namespace) for x in self.selection],
2277
+ hidden_components=self.hidden_components,
2278
+ local_concepts={
2279
+ x: y.with_namespace(namespace) for x, y in self.local_concepts.items()
2280
+ },
2281
+ order_by=self.order_by.with_namespace(namespace) if self.order_by else None,
2282
+ limit=self.limit,
2283
+ meta=self.meta,
2284
+ grain=self.grain.with_namespace(namespace),
2285
+ where_clause=(
2286
+ self.where_clause.with_namespace(namespace)
2287
+ if self.where_clause
2288
+ else None
2289
+ ),
2290
+ having_clause=(
2291
+ self.having_clause.with_namespace(namespace)
2292
+ if self.having_clause
2293
+ else None
2294
+ ),
2295
+ )
2296
+
2297
+
2298
+ class MultiSelectLineage(Mergeable, ConceptArgs, Namespaced, BaseModel):
2299
+ selects: List[SelectLineage]
2300
+ align: AlignClause
2301
+
2302
+ namespace: str
2303
+ order_by: Optional[OrderBy] = None
2304
+ limit: Optional[int] = None
2305
+ where_clause: Union["WhereClause", None] = Field(default=None)
2306
+ having_clause: Union["HavingClause", None] = Field(default=None)
2307
+ derive: DeriveClause | None = None
2308
+ hidden_components: set[str]
2309
+
2310
+ @property
2311
+ def grain(self):
2312
+ base = Grain()
2313
+ for select in self.selects:
2314
+ base += select.grain
2315
+ return base
2316
+
2317
+ @property
2318
+ def output_components(self) -> list[ConceptRef]:
2319
+ output = [
2320
+ ConceptRef.model_construct(address=x, datatype=DataType.UNKNOWN)
2321
+ for x in self.derived_concepts
2322
+ ]
2323
+ for select in self.selects:
2324
+ output += select.output_components
2325
+ return [x for x in output if x.address not in self.hidden_components]
2326
+
2327
+ def with_merge(
2328
+ self, source: Concept, target: Concept, modifiers: List[Modifier]
2329
+ ) -> MultiSelectLineage:
2330
+ new = MultiSelectLineage.model_construct(
2331
+ selects=[s.with_merge(source, target, modifiers) for s in self.selects],
2332
+ align=self.align,
2333
+ derive=(
2334
+ self.derive.with_merge(source, target, modifiers)
2335
+ if self.derive
2336
+ else None
2337
+ ),
2338
+ namespace=self.namespace,
2339
+ hidden_components=self.hidden_components,
2340
+ order_by=(
2341
+ self.order_by.with_merge(source, target, modifiers)
2342
+ if self.order_by
2343
+ else None
2344
+ ),
2345
+ limit=self.limit,
2346
+ where_clause=(
2347
+ self.where_clause.with_merge(source, target, modifiers)
2348
+ if self.where_clause
2349
+ else None
2350
+ ),
2351
+ having_clause=(
2352
+ self.having_clause.with_merge(source, target, modifiers)
2353
+ if self.having_clause
2354
+ else None
2355
+ ),
2356
+ )
2357
+ return new
2358
+
2359
+ def with_namespace(self, namespace: str) -> "MultiSelectLineage":
2360
+ return MultiSelectLineage.model_construct(
2361
+ selects=[c.with_namespace(namespace) for c in self.selects],
2362
+ align=self.align.with_namespace(namespace),
2363
+ derive=self.derive.with_namespace(namespace) if self.derive else None,
2364
+ namespace=namespace,
2365
+ hidden_components=self.hidden_components,
2366
+ order_by=self.order_by.with_namespace(namespace) if self.order_by else None,
2367
+ limit=self.limit,
2368
+ where_clause=(
2369
+ self.where_clause.with_namespace(namespace)
2370
+ if self.where_clause
2371
+ else None
2372
+ ),
2373
+ having_clause=(
2374
+ self.having_clause.with_namespace(namespace)
2375
+ if self.having_clause
2376
+ else None
2377
+ ),
2378
+ )
2379
+
2380
+ @property
2381
+ def derived_concepts(self) -> set[str]:
2382
+ output = set()
2383
+ for item in self.align.items:
2384
+ output.add(item.aligned_concept)
2385
+ if self.derive:
2386
+ for ditem in self.derive.items:
2387
+ output.add(ditem.derived_concept)
2388
+ return output
2389
+
2390
+ @property
2391
+ def concept_arguments(self):
2392
+ output = []
2393
+ for select in self.selects:
2394
+ output += select.output_components
2395
+ return unique(output, "address")
2396
+
2397
+
2398
+ class LooseConceptList(BaseModel):
2399
+ concepts: Sequence[Concept | ConceptRef]
2400
+
2401
+ @cached_property
2402
+ def addresses(self) -> set[str]:
2403
+ return {s.address for s in self.concepts}
2404
+
2405
+ @classmethod
2406
+ def validate(cls, v):
2407
+ return cls(v)
2408
+
2409
+ @cached_property
2410
+ def sorted_addresses(self) -> List[str]:
2411
+ return sorted(list(self.addresses))
2412
+
2413
+ def __str__(self) -> str:
2414
+ return f"lcl{str(self.sorted_addresses)}"
2415
+
2416
+ def __iter__(self):
2417
+ return iter(self.concepts)
2418
+
2419
+ def __eq__(self, other):
2420
+ if not isinstance(other, LooseConceptList):
2421
+ return False
2422
+ return self.addresses == other.addresses
2423
+
2424
+ def issubset(self, other):
2425
+ if not isinstance(other, LooseConceptList):
2426
+ return False
2427
+ return self.addresses.issubset(other.addresses)
2428
+
2429
+ def __contains__(self, other):
2430
+ if isinstance(other, str):
2431
+ return other in self.addresses
2432
+ if not isinstance(other, Concept):
2433
+ return False
2434
+ return other.address in self.addresses
2435
+
2436
+ def difference(self, other):
2437
+ if not isinstance(other, LooseConceptList):
2438
+ return False
2439
+ return self.addresses.difference(other.addresses)
2440
+
2441
+ def isdisjoint(self, other):
2442
+ if not isinstance(other, LooseConceptList):
2443
+ return False
2444
+ return self.addresses.isdisjoint(other.addresses)
2445
+
2446
+
2447
+ class AlignItem(Namespaced, BaseModel):
2448
+ alias: str
2449
+ concepts: List[ConceptRef]
2450
+ namespace: str = Field(default=DEFAULT_NAMESPACE, validate_default=True)
2451
+
2452
+ @field_validator("concepts", mode="before")
2453
+ @classmethod
2454
+ def enforce_concept_ref(cls, v):
2455
+ output = []
2456
+ for item in v:
2457
+ if isinstance(item, Concept):
2458
+ output.append(item.reference)
2459
+ else:
2460
+ output.append(item)
2461
+ return output
2462
+
2463
+ @computed_field # type: ignore
2464
+ @cached_property
2465
+ def concepts_lcl(self) -> LooseConceptList:
2466
+ return LooseConceptList(concepts=self.concepts)
2467
+
2468
+ @property
2469
+ def aligned_concept(self) -> str:
2470
+ return f"{self.namespace}.{self.alias}"
2471
+
2472
+ def with_namespace(self, namespace: str) -> "AlignItem":
2473
+ return AlignItem.model_construct(
2474
+ alias=self.alias,
2475
+ concepts=[c.with_namespace(namespace) for c in self.concepts],
2476
+ namespace=namespace,
2477
+ )
2478
+
2479
+
2480
+ class CustomFunctionFactory:
2481
+ def __init__(
2482
+ self,
2483
+ function: Expr,
2484
+ namespace: str,
2485
+ function_arguments: list[ArgBinding],
2486
+ name: str,
2487
+ ):
2488
+ self.namespace = namespace
2489
+ self.function = function
2490
+ self.function_arguments = function_arguments
2491
+ self.name = name
2492
+
2493
+ def with_namespace(self, namespace: str):
2494
+ self.namespace = namespace
2495
+ self.function = (
2496
+ self.function.with_namespace(namespace)
2497
+ if isinstance(self.function, Namespaced)
2498
+ else self.function
2499
+ )
2500
+ self.function_arguments = [
2501
+ x.with_namespace(namespace) for x in self.function_arguments
2502
+ ]
2503
+ return self
2504
+
2505
+ def __call__(self, *creation_args: ArgBinding | Expr):
2506
+ nout = (
2507
+ self.function.model_copy(deep=True)
2508
+ if isinstance(self.function, BaseModel)
2509
+ else self.function
2510
+ )
2511
+ creation_arg_list: list[ArgBinding | Expr] = list(creation_args)
2512
+ if len(creation_args) < len(self.function_arguments):
2513
+ for binding in self.function_arguments[len(creation_arg_list) :]:
2514
+ if binding.default is None:
2515
+ raise ValueError(f"Missing argument {binding.name}")
2516
+
2517
+ creation_arg_list.append(binding.default)
2518
+ for arg_idx, arg in enumerate(self.function_arguments):
2519
+ if not arg.datatype or arg.datatype == DataType.UNKNOWN:
2520
+ continue
2521
+ if arg_idx > len(creation_arg_list):
2522
+ continue
2523
+ comparison = arg_to_datatype(creation_arg_list[arg_idx])
2524
+ if comparison != arg.datatype:
2525
+ raise TypeError(
2526
+ f"Invalid type passed into custom function @{self.name} in position {arg_idx+1} for argument {arg.name}, expected {arg.datatype}, got {comparison}"
2527
+ )
2528
+ if isinstance(arg.datatype, TraitDataType):
2529
+ if not (
2530
+ isinstance(comparison, TraitDataType)
2531
+ and all(x in comparison.traits for x in arg.datatype.traits)
2532
+ ):
2533
+ raise TypeError(
2534
+ f"Invalid argument type passed into custom function @{self.name} in position {arg_idx+1} for argument {arg.name}, expected traits {arg.datatype.traits}, got {comparison}"
2535
+ )
2536
+
2537
+ if isinstance(nout, Mergeable):
2538
+ for idx, x in enumerate(creation_arg_list):
2539
+ if self.namespace == DEFAULT_NAMESPACE:
2540
+ target = f"{DEFAULT_NAMESPACE}.{self.function_arguments[idx].name}"
2541
+ else:
2542
+ target = self.function_arguments[idx].name
2543
+ nout = nout.with_reference_replacement(target, x)
2544
+ return nout
2545
+
2546
+
2547
+ class Metadata(BaseModel):
2548
+ """Metadata container object.
2549
+ TODO: support arbitrary tags"""
2550
+
2551
+ description: Optional[str] = None
2552
+ line_number: Optional[int] = None
2553
+ concept_source: ConceptSource = ConceptSource.MANUAL
2554
+
2555
+
2556
+ class Window(BaseModel):
2557
+ count: int
2558
+ window_order: WindowOrder
2559
+
2560
+ def __str__(self):
2561
+ return f"Window<{self.window_order}>"
2562
+
2563
+
2564
+ class WindowItemOver(BaseModel):
2565
+ contents: List[ConceptRef]
2566
+
2567
+
2568
+ class WindowItemOrder(BaseModel):
2569
+ contents: List["OrderItem"]
2570
+
2571
+
2572
+ class Comment(BaseModel):
2573
+ text: str
2574
+
2575
+
2576
+ class ArgBinding(Namespaced, DataTyped, BaseModel):
2577
+ name: str
2578
+ default: Expr | None = None
2579
+ datatype: (
2580
+ DataType | MapType | ArrayType | NumericType | StructType | TraitDataType
2581
+ ) = DataType.UNKNOWN
2582
+
2583
+ def with_namespace(self, namespace):
2584
+ return ArgBinding(
2585
+ name=address_with_namespace(self.name, namespace),
2586
+ default=(
2587
+ self.default.with_namespace(namespace)
2588
+ if isinstance(self.default, Namespaced)
2589
+ else self.default
2590
+ ),
2591
+ )
2592
+
2593
+ @property
2594
+ def output_datatype(self):
2595
+ if self.default is not None:
2596
+ return arg_to_datatype(self.default)
2597
+ return self.datatype
2598
+
2599
+
2600
+ class CustomType(BaseModel):
2601
+ name: str
2602
+ type: DataType | list[DataType]
2603
+ drop_on: list[FunctionType] = Field(default_factory=list)
2604
+ add_on: list[FunctionType] = Field(default_factory=list)
2605
+
2606
+ def with_namespace(self, namespace: str) -> "CustomType":
2607
+ return CustomType.model_construct(
2608
+ name=address_with_namespace(self.name, namespace),
2609
+ type=self.type,
2610
+ drop_on=self.drop_on,
2611
+ add_on=self.add_on,
2612
+ )
2613
+
2614
+
2615
+ Expr = (
2616
+ MagicConstants
2617
+ | bool
2618
+ | int
2619
+ | str
2620
+ | float
2621
+ | date
2622
+ | datetime
2623
+ | TupleWrapper
2624
+ | ListWrapper
2625
+ | MapWrapper
2626
+ | WindowItem
2627
+ | FilterItem
2628
+ | ConceptRef
2629
+ | Comparison
2630
+ | Conditional
2631
+ | FunctionCallWrapper
2632
+ | Parenthetical
2633
+ | Function
2634
+ | AggregateWrapper
2635
+ | CaseWhen
2636
+ | CaseElse
2637
+ )
2638
+
2639
+ FuncArgs = (
2640
+ ConceptRef
2641
+ | AggregateWrapper
2642
+ | Function
2643
+ | FunctionCallWrapper
2644
+ | Parenthetical
2645
+ | CaseWhen
2646
+ | CaseElse
2647
+ | WindowItem
2648
+ | FilterItem
2649
+ | bool
2650
+ | int
2651
+ | float
2652
+ | DatePart
2653
+ | str
2654
+ | date
2655
+ | datetime
2656
+ | MapWrapper[Any, Any]
2657
+ | TraitDataType
2658
+ | DataType
2659
+ | ArrayType
2660
+ | MapType
2661
+ | NumericType
2662
+ | ListWrapper[Any]
2663
+ | TupleWrapper[Any]
2664
+ | Comparison
2665
+ | Conditional
2666
+ | MagicConstants
2667
+ | ArgBinding
2668
+ | Ordering
2669
+ )