|
23 | 23 | from .data_structures.associative_array import AssociativeArray |
24 | 24 | from .data_structures.json_array import JSONArray |
25 | 25 | from .data_structures.key_value_pair import KeyValuePair |
| 26 | +from .data_structures.list_comprehension import ListComprehension |
26 | 27 | from .data_structures.lookup import Lookup |
27 | 28 | from .data_structures.range_lookup import RangeLookup |
28 | 29 | from .expressions.expression import Expression |
@@ -877,6 +878,13 @@ def _parse_operand(self, expression: Expression) -> bool: |
877 | 878 | lookup = self._parse_lookup(sub) |
878 | 879 | expression.add_node(lookup) |
879 | 880 | return True |
| 881 | + elif self.token.is_opening_bracket() and self._looks_like_list_comprehension(): |
| 882 | + list_comp = self._parse_list_comprehension() |
| 883 | + if list_comp is None: |
| 884 | + raise ValueError("Expected list comprehension") |
| 885 | + lookup = self._parse_lookup(list_comp) |
| 886 | + expression.add_node(lookup) |
| 887 | + return True |
880 | 888 | elif self.token.is_opening_brace() or self.token.is_opening_bracket(): |
881 | 889 | json = self._parse_json() |
882 | 890 | if json is None: |
@@ -1290,6 +1298,84 @@ def _parse_function_parameters(self) -> Iterator[ASTNode]: |
1290 | 1298 | break |
1291 | 1299 | self.set_next_token() |
1292 | 1300 |
|
| 1301 | + def _looks_like_list_comprehension(self) -> bool: |
| 1302 | + """Peek ahead from an opening bracket to determine whether the |
| 1303 | + upcoming tokens form a list comprehension (e.g. ``[n IN list | n.name]``) |
| 1304 | + rather than a plain JSON array literal (e.g. ``[1, 2, 3]``). |
| 1305 | +
|
| 1306 | + The heuristic is: ``[`` identifier ``IN`` -> list comprehension. |
| 1307 | + """ |
| 1308 | + saved_index = self._token_index |
| 1309 | + self.set_next_token() # skip '[' |
| 1310 | + self._skip_whitespace_and_comments() |
| 1311 | + |
| 1312 | + if not self.token.is_identifier_or_keyword(): |
| 1313 | + self._token_index = saved_index |
| 1314 | + return False |
| 1315 | + |
| 1316 | + self.set_next_token() # skip identifier |
| 1317 | + self._skip_whitespace_and_comments() |
| 1318 | + result = self.token.is_in() |
| 1319 | + self._token_index = saved_index |
| 1320 | + return result |
| 1321 | + |
| 1322 | + def _parse_list_comprehension(self) -> Optional[ListComprehension]: |
| 1323 | + """Parse a list comprehension expression. |
| 1324 | +
|
| 1325 | + Syntax: ``[variable IN list [WHERE condition] [| expression]]`` |
| 1326 | + """ |
| 1327 | + if not self.token.is_opening_bracket(): |
| 1328 | + return None |
| 1329 | + |
| 1330 | + list_comp = ListComprehension() |
| 1331 | + self.set_next_token() # skip '[' |
| 1332 | + self._skip_whitespace_and_comments() |
| 1333 | + |
| 1334 | + # Parse iteration variable |
| 1335 | + if not self.token.is_identifier_or_keyword(): |
| 1336 | + raise ValueError("Expected identifier") |
| 1337 | + reference = Reference(self.token.value or "") |
| 1338 | + self._state.variables[reference.identifier] = reference |
| 1339 | + list_comp.add_child(reference) |
| 1340 | + self.set_next_token() |
| 1341 | + self._expect_and_skip_whitespace_and_comments() |
| 1342 | + |
| 1343 | + # Parse IN keyword |
| 1344 | + if not self.token.is_in(): |
| 1345 | + raise ValueError("Expected IN") |
| 1346 | + self.set_next_token() |
| 1347 | + self._expect_and_skip_whitespace_and_comments() |
| 1348 | + |
| 1349 | + # Parse source array expression |
| 1350 | + array_expr = self._parse_expression() |
| 1351 | + if array_expr is None: |
| 1352 | + raise ValueError("Expected expression") |
| 1353 | + list_comp.add_child(array_expr) |
| 1354 | + |
| 1355 | + # Optional WHERE clause |
| 1356 | + self._skip_whitespace_and_comments() |
| 1357 | + where = self._parse_where() |
| 1358 | + if where is not None: |
| 1359 | + list_comp.add_child(where) |
| 1360 | + |
| 1361 | + # Optional | mapping expression |
| 1362 | + self._skip_whitespace_and_comments() |
| 1363 | + if self.token.is_pipe(): |
| 1364 | + self.set_next_token() |
| 1365 | + self._skip_whitespace_and_comments() |
| 1366 | + return_expr = self._parse_expression() |
| 1367 | + if return_expr is None: |
| 1368 | + raise ValueError("Expected expression after |") |
| 1369 | + list_comp.add_child(return_expr) |
| 1370 | + |
| 1371 | + self._skip_whitespace_and_comments() |
| 1372 | + if not self.token.is_closing_bracket(): |
| 1373 | + raise ValueError("Expected closing bracket") |
| 1374 | + self.set_next_token() |
| 1375 | + |
| 1376 | + del self._state.variables[reference.identifier] |
| 1377 | + return list_comp |
| 1378 | + |
1293 | 1379 | def _parse_json(self) -> Optional[ASTNode]: |
1294 | 1380 | if self.token.is_opening_brace(): |
1295 | 1381 | return self._parse_associative_array() |
|
0 commit comments