Codegen/Rust: allow breaking up schema file

misc/codegen/codegen.py

@@ -43,8 +43,8 @@ def _parse_args() -> argparse.Namespace:
                         "compute QL imports and in some comments and as root for relative paths provided as options. "
                         "If not provided it defaults to the directory of the configuration file, if any")
     path_arguments = [
-        p.add_argument("--schema", default="schema.py",
-                       help="input schema file (default %(default)s)"),
+        p.add_argument("--schema",
+                       help="input schema file (default schema.py)"),
         p.add_argument("--dbscheme",
                        help="output file for dbscheme generation, input file for trap generation"),
         p.add_argument("--ql-output",
@@ -87,6 +87,8 @@ def _parse_args() -> argparse.Namespace:
                     setattr(opts, flag, getattr(defaults, flag))
         if opts.root_dir is None:
             opts.root_dir = opts.configuration_file.parent
+    if opts.schema is None:
+        opts.schema = "schema.py"
     if not opts.generate:
         p.error("Nothing to do, specify --generate")
     # absolutize all paths

misc/codegen/loaders/schemaloader.py

@@ -1,4 +1,5 @@
 """ schema loader """
+import sys
 
 import inflection
 import typing
@@ -140,6 +141,8 @@ def load(m: types.ModuleType) -> schema.Schema:
             continue
         if name.startswith("__"):
             continue
+        if isinstance(data, types.ModuleType):
+            continue
         cls = _get_class(data)
         if classes and not cls.bases:
             raise schema.Error(
@@ -160,7 +163,10 @@ def load(m: types.ModuleType) -> schema.Schema:
 
 
 def load_file(path: pathlib.Path) -> schema.Schema:
-    spec = importlib.util.spec_from_file_location("schema", path)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
+    assert path.suffix in ("", ".py")
+    sys.path.insert(0, str(path.parent))
+    try:
+        module = importlib.import_module(path.with_suffix("").name)
+    finally:
+        sys.path.remove(str(path.parent))
     return load(module)

rust/BUILD.bazel

@@ -9,7 +9,7 @@ package(default_visibility = ["//rust:__subpackages__"])
 
 filegroup(
     name = "schema",
-    srcs = ["schema.py"],
+    srcs = glob(["schema/*.py"]),
 )
 
 filegroup(

rust/codegen.conf

@@ -1,4 +1,5 @@
 # configuration file for Rust code generation default options
+--schema=schema
 --generate=dbscheme,rusttest,ql,rust
 --dbscheme=ql/lib/rust.dbscheme
 --ql-output=ql/lib/codeql/rust/elements/internal/generated

rust/ql/lib/rust.dbscheme

@@ -117,7 +117,7 @@ locatable_locations(
 );
 
 
-// from schema.py
+// from schema
 
 @element =
   @locatable

rust/schema/__init__.py

@@ -0,0 +1,16 @@
+"""
+Schema description
+
+This file should be kept simple:
+* no flow control
+* no aliases
+* only class definitions with annotations and `include` calls
+
+For how documentation of generated QL code works, please read `misc/codegen/schema_documentation.md`.
+"""
+
+from .prelude import *
+from .ast import *
+
+include("../shared/tree-sitter-extractor/src/generator/prefix.dbscheme")
+include("prefix.dbscheme")

rust/schema/ast.py

@@ -1,60 +1,4 @@
-"""
-Schema description
-
-This file should be kept simple:
-* no flow control
-* no aliases
-* only class definitions with annotations and `include` calls
-
-For how documentation of generated QL code works, please read `misc/codegen/schema_documentation.md`.
-"""
-
-from misc.codegen.lib.schemadefs import *
-
-include("../shared/tree-sitter-extractor/src/generator/prefix.dbscheme")
-include("prefix.dbscheme")
-
-
-@qltest.skip
-class Element:
-    pass
-
-
-@qltest.skip
-class Locatable(Element):
-    pass
-
-
-@qltest.skip
-class AstNode(Locatable):
-    pass
-
-
-@qltest.skip
-class Unextracted(Element):
-    """
-    The base class marking everything that was not properly extracted for some reason, such as:
-    * syntax errors
-    * insufficient context information
-    * yet unimplemented parts of the extractor
-    """
-    pass
-
-
-@qltest.skip
-class Missing(Unextracted):
-    """
-    The base class marking errors during parsing or resolution.
-    """
-
-
-@qltest.skip
-class Unimplemented(Unextracted):
-    """
-    The base class for unimplemented nodes. This is used to mark nodes that are not yet extracted.
-    """
-    pass
-
+from .prelude import *
 
 class AssocItem(AstNode):
     pass

rust/schema/prelude.py

@@ -0,0 +1,41 @@
+from misc.codegen.lib.schemadefs import *
+
+@qltest.skip
+class Element:
+    pass
+
+
+@qltest.skip
+class Locatable(Element):
+    pass
+
+
+@qltest.skip
+class AstNode(Locatable):
+    pass
+
+
+@qltest.skip
+class Unextracted(Element):
+    """
+    The base class marking everything that was not properly extracted for some reason, such as:
+    * syntax errors
+    * insufficient context information
+    * yet unimplemented parts of the extractor
+    """
+    pass
+
+
+@qltest.skip
+class Missing(Unextracted):
+    """
+    The base class marking errors during parsing or resolution.
+    """
+
+
+@qltest.skip
+class Unimplemented(Unextracted):
+    """
+    The base class for unimplemented nodes. This is used to mark nodes that are not yet extracted.
+    """
+    pass