Typhon-Language 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

Typhon/SourceMap/ast_match_based_map.py

@@ -1,152 +1,199 @@
-import ast
-from .datatype import Range, Pos, RangeIntervalTree
-from ..Grammar.typhon_ast import get_pos_attributes_if_exists
-from ..Driver.debugging import debug_print, debug_verbose_print
-from ..SourceMap.ast_matching import match_ast
-
-
-class MatchBasedSourceMap:
-    def __init__(
-        self,
-        origin_to_unparsed: dict[ast.AST, ast.AST],
-        unparsed_to_origin: dict[ast.AST, ast.AST],
-        source_code: str,
-        source_file: str,
-    ):
-        self.origin_to_unparsed = origin_to_unparsed
-        self.unparsed_to_origin = unparsed_to_origin
-        self.origin_interval_tree = RangeIntervalTree[ast.AST]()
-        self.unparsed_interval_tree = RangeIntervalTree[ast.AST]()
-        self.source_code = source_code
-        self.source_file = source_file
-        self._setup_interval_trees()
-
-    def _setup_interval_trees(self):
-        for origin_node, unparsed_node in self.origin_to_unparsed.items():
-            origin_pos = get_pos_attributes_if_exists(origin_node)
-            if origin_pos is not None:
-                debug_verbose_print(
-                    f"Adding to origin interval tree:\n  range={Range.from_pos_attr_may_not_end(origin_pos)}\n    {ast.dump(origin_node)}\n    pos: {origin_pos}"
-                )
-                self.origin_interval_tree.add(
-                    Range.from_pos_attr_may_not_end(origin_pos), origin_node
-                )
-            unparsed_pos = get_pos_attributes_if_exists(unparsed_node)
-            if unparsed_pos is not None:
-                debug_verbose_print(
-                    f"Adding to unparsed interval tree: range={Range.from_pos_attr_may_not_end(unparsed_pos)} {ast.dump(unparsed_node)}"
-                )
-                self.unparsed_interval_tree.add(
-                    Range.from_pos_attr_may_not_end(unparsed_pos), unparsed_node
-                )
-
-    # Assume range in base_node, apply the offset of the range in base_node to result_node
-    # TODO: This is valid range conversion only for very simple cases.
-    def _apply_offset_range(
-        self,
-        range: Range,
-        base_node: ast.AST,
-        result_node: ast.AST,
-    ) -> Range | None:
-        base_pos_attr = get_pos_attributes_if_exists(base_node)
-        if base_pos_attr is None:
-            return None
-        base_range = Range.from_pos_attr_may_not_end(base_pos_attr)
-        result_pos_attr = get_pos_attributes_if_exists(result_node)
-        if result_pos_attr is None:
-            return None
-        result_range = Range.from_pos_attr_may_not_end(result_pos_attr)
-        range_offset = base_range.calc_offset(range)
-        debug_verbose_print(
-            f"Offsetting range:\n  base_range={base_range}\n  range={range}\n  offset={range_offset}\n  result_range={result_range}\n  apply_offset:{result_range.start.apply_offset(range_offset)}"
-        )
-        return result_range.start.apply_offset(range_offset)
-
-    def _range_to(
-        self,
-        range: Range,
-        interval_tree: RangeIntervalTree[ast.AST],
-        mapping: dict[ast.AST, ast.AST],
-    ) -> Range | None:
-        nodes = interval_tree.minimal_containers(range)
-        debug_verbose_print(f"Mapping range: {range} nodes: {nodes}")
-        if nodes:
-            if len(nodes) == 1:
-                # The canonical node for the range
-                node_range, node = nodes[0]
-                if node in mapping:
-                    debug_verbose_print(
-                        f"node is one: {ast.dump(node)}\n  mapping to {ast.dump(mapping[node])}\n  range: {range}\n  node_range: {node_range}"
-                    )
-                    if node_range == range:
-                        # If the range matches exactly, no need to apply offset
-                        return Range.from_ast_node(mapping[node])
-                    # Use offset mapping to precisely map the range inside the node
-                    return self._apply_offset_range(range, node, mapping[node])
-            else:
-                # When multiple nodes are found, merge the ranges of all mapped nodes
-                mapped_nodes = (mapping[node] for _, node in nodes if node in mapping)
-                pos_attrs = (
-                    pos_attr
-                    for mapped_node in mapped_nodes
-                    if (pos_attr := get_pos_attributes_if_exists(mapped_node))
-                )
-                debug_verbose_print(f"Found {len(nodes)} nodes for the given range.")
-                return Range.merge_ranges(
-                    Range.from_pos_attr_may_not_end(pos_attr) for pos_attr in pos_attrs
-                )
-        debug_verbose_print("No nodes found for the given range.")
-
-    def unparsed_range_to_origin(
-        self,
-        range_unparsed: Range,
-    ) -> Range | None:
-        return self._range_to(
-            range_unparsed,
-            self.unparsed_interval_tree,
-            self.unparsed_to_origin,
-        )
-
-    def unparsed_range_to_source_code(
-        self,
-        range_unparsed: Range,
-    ) -> str | None:
-        range_in_origin = self.unparsed_range_to_origin(range_unparsed)
-        if range_in_origin is None:
-            return None
-        return range_in_origin.of_string(self.source_code)
-
-    def origin_range_to_unparsed(
-        self,
-        range_origin: Range,
-    ) -> Range | None:
-        debug_verbose_print(f"Mapping origin range: {range_origin}")
-        return self._range_to(
-            range_origin,
-            self.origin_interval_tree,
-            self.origin_to_unparsed,
-        )
-
-    def origin_node_to_unparsed_range(
-        self,
-        origin_node: ast.AST,
-    ) -> Range | None:
-        range_origin = get_pos_attributes_if_exists(origin_node)
-        if range_origin is None:
-            return None
-        range_origin_part = Range.from_pos_attr_may_not_end(range_origin)
-        return self.origin_range_to_unparsed(range_origin_part)
-
-
-def map_from_transformed_ast(
-    origin_ast: ast.AST,
-    unparsed_ast: ast.AST,
-    source_code: str,
-    source_file: str,
-) -> MatchBasedSourceMap | None:
-    mapping = match_ast(origin_ast, unparsed_ast)
-    if mapping is None:
-        return None
-    return MatchBasedSourceMap(
-        mapping.left_to_right, mapping.right_to_left, source_code, source_file
-    )
+import ast
+from .datatype import Range, Pos, RangeIntervalTree, RangeInterval
+from ..Grammar.typhon_ast import get_pos_attributes_if_exists
+from ..Driver.debugging import debug_print, debug_verbose_print
+from ..SourceMap.ast_matching import match_ast
+from .defined_name_retrieve import defined_name_retrieve
+
+
+class MatchBasedSourceMap:
+    def __init__(
+        self,
+        origin_to_unparsed: dict[ast.AST, ast.AST],
+        unparsed_to_origin: dict[ast.AST, ast.AST],
+        source_code: str,
+        source_file: str,
+        unparsed_code: str,
+    ):
+        self.origin_to_unparsed = origin_to_unparsed
+        self.unparsed_to_origin = unparsed_to_origin
+        self.origin_interval_tree = RangeIntervalTree[ast.AST]()
+        self.unparsed_interval_tree = RangeIntervalTree[ast.AST]()
+        self.source_code = source_code
+        self.source_file = source_file
+        self.unparsed_code: str = unparsed_code
+        self._setup_interval_trees()
+
+    def _setup_interval_trees(self):
+        for origin_node, unparsed_node in self.origin_to_unparsed.items():
+            origin_pos = get_pos_attributes_if_exists(origin_node)
+            if origin_pos is not None:
+                debug_verbose_print(
+                    f"Adding to origin interval tree:\n    range={Range.from_pos_attr_may_not_end(origin_pos)}\n    {ast.dump(origin_node)}\n    pos: {origin_pos}"
+                )
+                self.origin_interval_tree.add(
+                    Range.from_pos_attr_may_not_end(origin_pos), origin_node
+                )
+            unparsed_pos = get_pos_attributes_if_exists(unparsed_node)
+            if unparsed_pos is not None:
+                debug_verbose_print(
+                    f"  Adding to unparsed interval tree:\n    range={Range.from_pos_attr_may_not_end(unparsed_pos)}\n    {ast.dump(unparsed_node)}\n    pos: {unparsed_pos}"
+                )
+                self.unparsed_interval_tree.add(
+                    Range.from_pos_attr_may_not_end(unparsed_pos), unparsed_node
+                )
+
+    # Assume range in base_node, apply the offset of the range in base_node to result_node
+    # TODO: This is valid range conversion only for very simple cases.
+    def _apply_offset_range(
+        self,
+        range: Range,
+        base_node: ast.AST,
+        result_node: ast.AST,
+    ) -> Range | None:
+        base_pos_attr = get_pos_attributes_if_exists(base_node)
+        if base_pos_attr is None:
+            return None
+        base_range = Range.from_pos_attr_may_not_end(base_pos_attr)
+        result_pos_attr = get_pos_attributes_if_exists(result_node)
+        if result_pos_attr is None:
+            return None
+        result_range = Range.from_pos_attr_may_not_end(result_pos_attr)
+        range_offset = base_range.calc_offset(range)
+        debug_verbose_print(
+            f"Offsetting range:\n  base_range={base_range}\n  range={range}\n  offset={range_offset}\n  result_range={result_range}\n  apply_offset:{result_range.start.apply_offset(range_offset)}"
+        )
+        return result_range.start.apply_offset(range_offset)
+
+    def _range_to(
+        self,
+        range: Range,
+        interval_tree: RangeIntervalTree[ast.AST],
+        mapping: dict[ast.AST, ast.AST],
+    ) -> Range | None:
+        nodes = interval_tree.minimal_containers(range)
+        debug_verbose_print(f"Mapping range: {range} nodes: {nodes}")
+        if nodes:
+            if len(nodes) == 1:
+                # The canonical node for the range
+                node_range, node = nodes[0]
+                if node in mapping:
+                    debug_verbose_print(
+                        f"node is one: {ast.dump(node)}\n  mapping to {ast.dump(mapping[node])}\n  range: {range}\n  node_range: {node_range}"
+                    )
+                    if node_range == range:
+                        # If the range matches exactly, no need to apply offset
+                        return Range.from_ast_node(mapping[node])
+                    # Use offset mapping to precisely map the range inside the node
+                    return self._apply_offset_range(range, node, mapping[node])
+            else:
+                # When multiple nodes are found, merge the ranges of all mapped nodes
+                mapped_nodes = (mapping[node] for _, node in nodes if node in mapping)
+                pos_attrs = (
+                    pos_attr
+                    for mapped_node in mapped_nodes
+                    if (pos_attr := get_pos_attributes_if_exists(mapped_node))
+                )
+                debug_verbose_print(f"Found {len(nodes)} nodes for the given range.")
+                return Range.merge_ranges(
+                    Range.from_pos_attr_may_not_end(pos_attr) for pos_attr in pos_attrs
+                )
+        debug_verbose_print("No nodes found for the given range.")
+
+    def unparsed_range_to_origin(
+        self,
+        range_unparsed: Range,
+    ) -> Range | None:
+        return self._range_to(
+            range_unparsed,
+            self.unparsed_interval_tree,
+            self.unparsed_to_origin,
+        )
+
+    def unparsed_range_to_source_code(
+        self,
+        range_unparsed: Range,
+    ) -> str | None:
+        range_in_origin = self.unparsed_range_to_origin(range_unparsed)
+        if range_in_origin is None:
+            return None
+        return range_in_origin.of_string(self.source_code)
+
+    def unparsed_range_to_origin_node(
+        self,
+        range_unparsed: Range,
+        filter_node_type: type[ast.AST] | None = None,
+    ) -> ast.AST | None:
+        nodes: list[RangeInterval[ast.AST]] = (
+            self.unparsed_interval_tree.minimal_containers(range_unparsed)
+        )
+        debug_verbose_print(
+            f"Mapping unparsed range to origin node: {range_unparsed} nodes: {nodes}"
+        )
+        if filter_node_type is not None:
+            nodes = [
+                (r, n)
+                for r, n in nodes
+                if isinstance(self.unparsed_to_origin.get(n, None), filter_node_type)
+            ]
+        if nodes and len(nodes) == 1:
+            _, node = nodes[0]
+            return self.unparsed_to_origin.get(node, None)
+        debug_verbose_print("No nodes found for the given unparsed range.")
+        return None
+
+    def origin_range_to_unparsed(
+        self,
+        range_origin: Range,
+    ) -> Range | None:
+        debug_verbose_print(f"Mapping origin range: {range_origin}")
+        return self._range_to(
+            range_origin,
+            self.origin_interval_tree,
+            self.origin_to_unparsed,
+        )
+
+    def origin_node_to_unparsed_range(
+        self,
+        origin_node: ast.AST,
+    ) -> Range | None:
+        range_origin = get_pos_attributes_if_exists(origin_node)
+        if range_origin is None:
+            return None
+        range_origin_part = Range.from_pos_attr_may_not_end(range_origin)
+        return self.origin_range_to_unparsed(range_origin_part)
+
+
+def map_from_translated_ast(
+    origin_ast: ast.AST,
+    unparsed_ast: ast.AST,
+    source_code: str,
+    source_file_path: str,
+    unparsed_code: str,
+) -> MatchBasedSourceMap | None:
+    defined_name_retrieve(unparsed_ast, unparsed_code)
+    mapping = match_ast(origin_ast, unparsed_ast)
+    if mapping is None:
+        return None
+    return MatchBasedSourceMap(
+        mapping.left_to_right,
+        mapping.right_to_left,
+        source_code,
+        source_file_path,
+        unparsed_code,
+    )
+
+
+def map_from_translated(
+    origin_ast: ast.AST,
+    source_code: str,
+    source_file_path: str,
+    translated_code: str,
+) -> MatchBasedSourceMap | None:
+    return map_from_translated_ast(
+        origin_ast,
+        ast.parse(translated_code),
+        source_code,
+        source_file_path,
+        translated_code,
+    )

Typhon/SourceMap/ast_matching.py

@@ -1,87 +1,102 @@
-from dataclasses import dataclass
-import ast
-from contextlib import contextmanager
-from typing import Any, cast
-from ..Driver.debugging import debug_verbose_print
-
-
-# Match the AST node to right module recursively
-class MatchingVisitor(ast.NodeVisitor):
-    def __init__(
-        self,
-        right: ast.AST,
-        left_to_right: dict[ast.AST, ast.AST],
-        right_to_left: dict[ast.AST, ast.AST],
-    ):
-        self.left_to_right = left_to_right
-        self.right_to_left = right_to_left
-        self.right = right
-
-    @contextmanager
-    def _with_right(self, right: ast.AST):
-        old_right = self.right
-        self.right = right
-        yield
-        self.right = old_right
-
-    def _commit(self, left: ast.AST, right: ast.AST):
-        debug_verbose_print(f"Matched: {ast.dump(left)}  <->  {ast.dump(right)}")
-        self.left_to_right[left] = right
-        self.right_to_left[right] = left
-
-    def _visit_list(self, lefts: list[Any], rights: list[Any]):
-        if len(lefts) != len(rights):
-            # TODO: Error rescue: List length mismatch
-            raise ValueError(f"List length mismatch: {len(lefts)} vs {len(rights)}")
-        for left, right in zip(lefts, rights):
-            if not isinstance(left, ast.AST) or not isinstance(right, ast.AST):
-                if left != right:
-                    # TODO: Error rescue: List value mismatch
-                    raise ValueError(f"List value mismatch: {left} vs {right}")
-                continue
-            with self._with_right(right):
-                self.visit(left)
-
-    def visit(self, node: ast.AST):
-        right = self.right
-        if type(node) is not type(right):
-            # TODO: Error rescue: Type mismatch
-            raise ValueError(f"Type mismatch: {ast.dump(node)} vs {ast.dump(right)}")
-
-        self._commit(node, right)
-
-        for field, value in ast.iter_fields(node):
-            right_value = getattr(right, field, None)
-            if isinstance(value, list):
-                if isinstance(right_value, list):
-                    self._visit_list(
-                        cast(list[Any], value), cast(list[Any], right_value)
-                    )
-                else:
-                    # TODO: Error rescue: List length mismatch
-                    raise ValueError(
-                        f"List mismatch in field {field}: {ast.dump(node)} vs {ast.dump(right)}"
-                    )
-            elif isinstance(value, ast.AST):
-                with self._with_right(cast(ast.AST, right_value)):
-                    self.visit(value)
-            else:
-                if value != right_value:
-                    # TODO: Error rescue: Value mismatch
-                    raise ValueError(
-                        f"Value mismatch in field {field}: {value} vs {right_value}"
-                    )
-
-
-@dataclass
-class MatchResult:
-    left_to_right: dict[ast.AST, ast.AST]
-    right_to_left: dict[ast.AST, ast.AST]
-
-
-def match_ast[T: ast.AST](left: T, right: T) -> MatchResult | None:
-    left_to_right: dict[ast.AST, ast.AST] = {}
-    right_to_left: dict[ast.AST, ast.AST] = {}
-    visitor = MatchingVisitor(right, left_to_right, right_to_left)
-    visitor.visit(left)
-    return MatchResult(left_to_right, right_to_left)
+from dataclasses import dataclass
+import ast
+from contextlib import contextmanager
+from typing import Any, cast
+from ..Driver.debugging import debug_verbose_print
+from ..Grammar.typhon_ast import DefinesName, get_defined_name, get_import_from_names
+
+
+# Match the AST node to right module recursively
+class MatchingVisitor(ast.NodeVisitor):
+    def __init__(
+        self,
+        right: ast.AST,
+        left_to_right: dict[ast.AST, ast.AST],
+        right_to_left: dict[ast.AST, ast.AST],
+    ):
+        self.left_to_right = left_to_right
+        self.right_to_left = right_to_left
+        self.right = right
+
+    @contextmanager
+    def _with_right(self, right: ast.AST):
+        old_right = self.right
+        self.right = right
+        yield
+        self.right = old_right
+
+    def _commit(self, left: ast.AST, right: ast.AST):
+        debug_verbose_print(f"Matched: {ast.dump(left)}  <->  {ast.dump(right)}")
+        self.left_to_right[left] = right
+        self.right_to_left[right] = left
+
+    def _visit_list(self, lefts: list[Any], rights: list[Any]):
+        if len(lefts) != len(rights):
+            # TODO: Error rescue: List length mismatch
+            raise ValueError(f"List length mismatch: {len(lefts)} vs {len(rights)}")
+        for left, right in zip(lefts, rights):
+            if not isinstance(left, ast.AST) or not isinstance(right, ast.AST):
+                if left != right:
+                    # TODO: Error rescue: List value mismatch
+                    raise ValueError(f"List value mismatch: {left} vs {right}")
+                continue
+            with self._with_right(right):
+                self.visit(left)
+
+    def visit(self, node: ast.AST):
+        right = self.right
+        if type(node) is not type(right):
+            # TODO: Error rescue: Type mismatch
+            raise ValueError(f"Type mismatch: {ast.dump(node)} vs {ast.dump(right)}")
+        # Commit the match
+        self._commit(node, right)
+        # Check defined name
+        if isinstance(node, DefinesName):
+            left_name = get_defined_name(node)
+            right_name = get_defined_name(cast(DefinesName, right))
+            if left_name is not None and right_name is not None:
+                with self._with_right(right_name):
+                    self.visit(left_name)
+            # Allow defined name not matching
+        # Check import from module names
+        if isinstance(node, ast.ImportFrom):
+            modules = get_import_from_names(node)
+            right_modules = get_import_from_names(cast(ast.ImportFrom, right))
+            if modules and right_modules:
+                self._visit_list(modules, right_modules)
+        # Recursively visit fields
+        for field, value in ast.iter_fields(node):
+            right_value = getattr(right, field, None)
+            if isinstance(value, list):
+                if isinstance(right_value, list):
+                    self._visit_list(
+                        cast(list[Any], value), cast(list[Any], right_value)
+                    )
+                else:
+                    # TODO: Error rescue: List length mismatch
+                    raise ValueError(
+                        f"List mismatch in field {field}: {ast.dump(node)} vs {ast.dump(right)}"
+                    )
+            elif isinstance(value, ast.AST):
+                with self._with_right(cast(ast.AST, right_value)):
+                    self.visit(value)
+            else:
+                if value != right_value:
+                    # TODO: Error rescue: Value mismatch
+                    raise ValueError(
+                        f"Value mismatch in field {field}: {value} vs {right_value}"
+                    )
+
+
+@dataclass
+class MatchResult:
+    left_to_right: dict[ast.AST, ast.AST]
+    right_to_left: dict[ast.AST, ast.AST]
+
+
+def match_ast[T: ast.AST](left: T, right: T) -> MatchResult | None:
+    left_to_right: dict[ast.AST, ast.AST] = {}
+    right_to_left: dict[ast.AST, ast.AST] = {}
+    visitor = MatchingVisitor(right, left_to_right, right_to_left)
+    visitor.visit(left)
+    return MatchResult(left_to_right, right_to_left)