add custom symbol table and check/transform omp directive calls

cesarpomar · cesarpomar · commit 06c297aa06f0 · 2026-01-27T13:12:04.000+01:00
diff --git a/omp4py/core/preprocessor/obj2ast.py b/omp4py/core/preprocessor/obj2ast.py
@@ -23,73 +23,7 @@ def get_indentation(lines: list[str]) -> int:
     return 0
 
 
-class NamespaceVisitor(ast.NodeVisitor):
-    namespace: int
-    line: int
-    name: str
-    result: tuple[ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef, int] | None
-
-    def __init__(self, name: str, line: int) -> None:
-        self.namespace = 0
-        self.name = name
-        self.line: int = line
-        self.result = None
-
-    def search(self, module: ast.Module) -> tuple[ast.stmt, int]:
-        self.generic_visit(module)
-        if self.result is None:
-            msg: str = "source code not found"
-            raise ValueError(msg)
-        return self.result
-
-    def visit_namespace(self, node: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef) -> None:
-        self.namespace += 1
-        lineno: int = node.lineno - 1
-        if len(node.decorator_list) > 0:
-            lineno = node.decorator_list[0].lineno - 1
-        if lineno == self.line and node.name == self.name:
-            self.result = (node, self.namespace)
-
-    visit_FunctionDef = visit_namespace  # noqa: N815
-    visit_ClassDef = visit_namespace  # noqa: N815
-    visit_AsyncFunctionDef = visit_namespace  # noqa: N815
-
-
-def check_object(obj: Callable[..., Any] | type):
-    if not (inspect.isclass(obj) or inspect.iscoroutinefunction(obj) or inspect.isfunction(obj)):
-        msg: str = "Invalid object type: only classes, functions, or async functions are allowed"
-        raise ValueError(msg)
-    qname: str = getattr(obj, "__qualname__", "")
-
-    if "." in qname:
-        msg: str = "Decorator must be applied to an outer function or class"
-        filename: str = inspect.getfile(obj)
-        lines: list[str]
-        start: int
-        lines, start = inspect.findsource(obj)
-        offset: int = len(lines[start]) - len(lines[start].lstrip())
-        full_source: str = "".join(lines)
-        raise syntax_error(msg, Span(start + 1, offset), full_source, filename)
-
-
-def from_object(obj: Callable[..., Any] | type) -> tuple[str, str, ast.Module, int]:
-    check_object(obj)
-    lines: list[str]
-    start: int
-    lines, start = inspect.findsource(obj)
-    full_source: str = "".join(lines)
-    filename: str = inspect.getfile(obj)
-    module: ast.Module = ast.parse(full_source, filename)
-    name: str = getattr(obj, "__name__", "")
-    obj_smt: ast.stmt
-    namespace: int
-    obj_smt, namespace = NamespaceVisitor(name, start).search(module)
-    module.body = [obj_smt]
-
-    return filename, full_source, module, namespace
-
-
-def from_object2(obj: Callable[..., Any] | type) -> tuple[str, str, ast.Module]:
+def from_object(obj: Callable[..., Any] | type) -> tuple[str, str, ast.Module]:
     lines: list[str]
     start: int
     lines, start = inspect.findsource(obj)
diff --git a/omp4py/core/preprocessor/preprocessor.py b/omp4py/core/preprocessor/preprocessor.py
@@ -23,9 +23,8 @@ def process_object[T: Callable[..., Any] | type](arg: T, params: Params) -> T:
     filename: str
     data: str
     module: ast.Module
-    namespace: int
-    filename, data, module,namespace = obj2ast.from_object(arg)
-    module: ast.Module = process(module, namespace,data, filename, params)
+    filename, data, module = obj2ast.from_object(arg)
+    module: ast.Module = process(module, False,data, filename, params)
     globals: ModuleType = sys.modules[arg.__module__]
 
     # TODO: use importlib to use pyc cache
@@ -41,7 +40,7 @@ def process_object[T: Callable[..., Any] | type](arg: T, params: Params) -> T:
 
 
 def process_source(data: str, filename: str, params: Params) -> ast.Module:
-    return process(ast.parse(data, filename),0, data, filename, params)
+    return process(ast.parse(data, filename),True, data, filename, params)
 
 
 def process_file(filename: str, params: Params) -> str:
@@ -58,7 +57,7 @@ def process_file(filename: str, params: Params) -> str:
     return filename
 
 
-def process(module: ast.Module, namespace: int, full_source: str, filename: str, params: Params) -> ast.Module:
-    transformer: OmpTransformer = OmpTransformer(full_source, filename, module, namespace, params)
+def process(module: ast.Module, is_module: bool, full_source: str, filename: str, params: Params) -> ast.Module:
+    transformer: OmpTransformer = OmpTransformer(full_source, filename, module, is_module, params)
 
     return transformer.transform()
diff --git a/omp4py/core/preprocessor/transformers/context.py b/omp4py/core/preprocessor/transformers/context.py
@@ -1,10 +1,16 @@
-from omp4py.core.directive.schema import Directive
 import ast
 import os
-import symtable
 from dataclasses import dataclass, field
+from symtable import symtable as native_symtable
 
-__all__ = ["Context", "Params"]
+from omp4py.core.parser import Directive
+from omp4py.core.preprocessor.transformers.symtable import SymbolTable
+
+__all__ = ["Context", "Params", "SymbolTable", "global_symtable"]
+
+
+def global_symtable(data: str, filename: str) -> "SymbolTable":
+    return SymbolTable(list(native_symtable(data, filename, "exec").get_identifiers()))
 
 
 def environ_bool(key: str, default: bool) -> bool:
@@ -24,19 +30,25 @@ class Context:
     full_source: str
     module: ast.Module
     namespace: int
-    symtable_stack: list[symtable.SymbolTable]
+    symtable: SymbolTable
     node_stack: list[ast.AST]
     decorator: ast.expr | None
     directive: Directive | None
 
-    def __init__(self, full_source: str, filename: str, module: ast.Module, namespace: int, params: Params) -> None:
+    def __init__(
+        self,
+        full_source: str,
+        filename: str,
+        module: ast.Module,
+        is_module: bool,
+        params: Params,
+    ) -> None:
         self.params = params
         self.filename = filename
         self.module = module
-        self.namespace = namespace
+        self.is_module = is_module
         self.full_source = full_source
-        self.symtable_stack = []
+        self.symtable = SymbolTable()
         self.node_stack = [module]
         self.decorator = None
         self.directive = None
-
diff --git a/omp4py/core/preprocessor/transformers/symtable.py b/omp4py/core/preprocessor/transformers/symtable.py
@@ -0,0 +1,166 @@
+import ast
+import re
+from collections.abc import Callable, KeysView
+from dataclasses import dataclass
+from typing import Optional
+
+PREFIX: str = "_omp_"
+
+var_check: re.Pattern[str] = re.compile(rf"^({PREFIX})?([0-9]+)?(.+)$")
+
+
+def is_variable(name: str) -> bool:
+    return not name.startswith(PREFIX) or name[len(PREFIX) : len(PREFIX) + 1].isdigit()
+
+
+@dataclass
+class SymbolEntry:
+    scope_name: str
+    old_name: str
+    used: bool = False
+    assigned: bool = False
+    annotation: ast.expr | None = None
+
+
+class SymbolTableVisitor(ast.NodeVisitor):
+    symbols: dict[str, SymbolEntry]
+    check_namespace: bool
+    to_rename: set[str]
+
+    def __init__(self, global_vars: list[str] | None):
+        self.to_rename = set()
+        self.check_namespace = False
+        self.symbols = {}
+        if global_vars is not None:
+            name: str
+            for name in global_vars:
+                self.update_symbol(name, lambda x: None).assigned = True
+
+    def update_symbol(self, name: str, rename: Callable[[str], None]) -> SymbolEntry:
+        match: re.Match[str] | None = var_check.match(name)
+        if not match:
+            return SymbolEntry("", "")
+        omp: str | None = match.group(1)
+        n: str | None = match.group(2)
+        real_name: str = match.group(3)
+        if omp is not None and n is None:
+            return SymbolEntry("", "")  # internal _omp_ var
+
+        symbol: SymbolEntry
+        if real_name in self.symbols:
+            symbol = self.symbols[real_name]
+        else:
+            old_name: str = real_name if n is None else (PREFIX + str(int(n) - 1) + real_name)
+            symbol = self.symbols[real_name] = SymbolEntry(name, old_name)
+
+        if real_name in self.to_rename:
+            new_name: str = PREFIX + str(1 if n is None else (int(n) + 1)) + real_name
+            symbol.old_name = symbol.scope_name
+            symbol.scope_name = new_name
+            rename(new_name)
+        return symbol
+
+    def check(self, node: ast.AST, namespace: bool = False) -> None:
+        self.check_namespace = False
+        self.visit(node)
+        if namespace:
+            self.check_namespace = namespace
+            super().generic_visit(node)
+
+    def rename(self, names: set[str], node: ast.AST) -> None:
+        self.to_rename = names
+        self.visit(node)
+
+    def visit(self, node: ast.AST) -> None:
+        super().visit(node)
+        if len(self.to_rename) > 0 or (
+            self.check_namespace and not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef))
+        ):
+            super().generic_visit(node)
+
+    def generic_visit(self, node: ast.AST):
+        pass
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        self.update_symbol(node.name, lambda x: setattr(node, "name", x)).assigned = True
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
+        self.update_symbol(node.name, lambda x: setattr(node, "name", x)).assigned = True
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        self.update_symbol(node.name, lambda x: setattr(node, "name", x)).assigned = True
+
+    def visit_Import(self, node: ast.Import) -> None:
+        alias: ast.alias
+        for alias in node.names:
+            name: str = alias.name if alias.asname is None else alias.asname
+            self.update_symbol(name, lambda x: setattr(node, "asname", x)).assigned = True
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        alias: ast.alias
+        for alias in node.names:
+            name: str = alias.name if alias.asname is None else alias.asname
+            self.update_symbol(name, lambda x: setattr(alias, "asname", x)).assigned = True  # noqa: B023
+
+    def visit_Global(self, node: ast.Global) -> None:
+        i: int
+        name: str
+        for i, name in enumerate(node.names):
+            self.update_symbol(name, lambda x: node.names.__setitem__(i, x)).assigned = True  # noqa: B023
+
+    def visit_Nonlocal(self, node: ast.Nonlocal) -> None:
+        i: int
+        name: str
+        for i, name in enumerate(node.names):
+            self.update_symbol(name, lambda x: node.names.__setitem__(i, x)).assigned = True  # noqa: B023
+
+    def visit_Name(self, node: ast.Name) -> None:
+        if isinstance(node.ctx, ast.Load):
+            self.update_symbol(node.id, lambda x: setattr(node, "id", x)).used = True
+        else:
+            self.update_symbol(node.id, lambda x: setattr(node, "id", x)).assigned = True
+
+    def visit_arg(self, node: ast.arg) -> None:
+        s: SymbolEntry = self.update_symbol(node.arg, lambda x: setattr(node, "arg", x))
+        s.assigned = True
+        s.annotation = node.annotation
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        if isinstance(node.target, ast.Name):
+            self.update_symbol(node.target.id, lambda x: setattr(node.target, "id", x)).annotation = node.annotation
+
+
+class SymbolTable:
+    _parent: Optional["SymbolTable"]
+    _visitor: SymbolTableVisitor
+
+    def __init__(self, global_vars: list[str] | None = None):
+        self._visitor = SymbolTableVisitor(global_vars)
+
+    def check_namespace(self, node: ast.AST) -> "SymbolTable":
+        child: SymbolTable = self.new_child()
+        child._visitor.check(node, namespace=True)
+        return child
+
+    def update(self, node: ast.AST) -> None:
+        self._visitor.check(node)
+
+    def new_child(self) -> "SymbolTable":
+        child: SymbolTable = SymbolTable()
+        child._parent = self
+        return child
+
+    def rename(self, names: set[str], node: ast.AST) -> None:
+        self._visitor.rename(names, node)
+
+    def symbols(self) -> list[SymbolEntry]:
+        return list(self._visitor.symbols.values())
+
+    def identifiers(self) -> KeysView[str]:
+        return self._visitor.symbols.keys()
+
+    def __getattr__(self, name: str) -> SymbolEntry:
+        return self._visitor.symbols[name]
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._visitor.symbols
diff --git a/omp4py/core/preprocessor/transformers/transformer.py b/omp4py/core/preprocessor/transformers/transformer.py
