diff --git a/src/fpy/semantics.py b/src/fpy/semantics.py index b25a2a8..1ecf1a5 100644 --- a/src/fpy/semantics.py +++ b/src/fpy/semantics.py @@ -1,9 +1,8 @@ from __future__ import annotations -from dataclasses import fields, replace as dc_replace +from dataclasses import replace as dc_replace from datetime import datetime, timezone from decimal import Decimal import decimal -import itertools import math from pathlib import Path import struct @@ -68,6 +67,7 @@ is_symbol_an_expr, ) from fpy.visitors import ( + ast_children, STOP_DESCENT, TopDownVisitor, Visitor, @@ -139,19 +139,7 @@ def run(self, start: Ast, state: CompileState): def _descend(node: Ast): if not isinstance(node, Ast): return - children = [] - for field in fields(node): - field_val = getattr(node, field.name) - if isinstance(field_val, list): - if len(field_val) > 0 and isinstance(field_val[0], tuple): - field_val = itertools.chain.from_iterable(field_val) - children.extend(field_val) - else: - children.append(field_val) - - for child in children: - if not isinstance(child, Ast): - continue + for child in ast_children(node): self._visit(child, state) state.parent_map[child] = node if len(state.errors) != 0: diff --git a/src/fpy/visitors.py b/src/fpy/visitors.py index 8bcacbe..abb4c55 100644 --- a/src/fpy/visitors.py +++ b/src/fpy/visitors.py @@ -30,6 +30,21 @@ def __repr__(self): STOP_DESCENT = _StopDescent() +def ast_children(node: Ast) -> list[Ast]: + """Snapshot a node's AST children in dataclass field order.""" + children = [] + for field in fields(node): + value = getattr(node, field.name) + if isinstance(value, list): + # Function parameters are stored as a list of tuples. + if value and isinstance(value[0], tuple): + value = itertools.chain.from_iterable(value) + children.extend(value) + else: + children.append(value) + return [child for child in children if isinstance(child, Ast)] + + class Visitor: """visits each class, calling a custom visit function, if one is defined, for each node type""" @@ -87,20 +102,7 @@ def run(self, start: Ast, state: CompileState): def _descend(node: Ast): if not isinstance(node, Ast): return - children = [] - for field in fields(node): - field_val = getattr(node, field.name) - if isinstance(field_val, list): - # also handle the one case where we have a list of tuples - if len(field_val) > 0 and isinstance(field_val[0], tuple): - field_val = itertools.chain.from_iterable(field_val) - children.extend(field_val) - else: - children.append(field_val) - - for child in children: - if not isinstance(child, Ast): - continue + for child in ast_children(node): _descend(child) if len(state.errors) != 0: break @@ -121,20 +123,7 @@ def run(self, start: Ast, state: CompileState): def _descend(node: Ast): if not isinstance(node, Ast): return - children = [] - for field in fields(node): - field_val = getattr(node, field.name) - if isinstance(field_val, list): - # also handle the one case where we have a list of tuples - if len(field_val) > 0 and isinstance(field_val[0], tuple): - field_val = itertools.chain.from_iterable(field_val) - children.extend(field_val) - else: - children.append(field_val) - - for child in children: - if not isinstance(child, Ast): - continue + for child in ast_children(node): result = self._visit(child, state) if len(state.errors) != 0: break