Codegen: allow full annotation of classes

misc/codegen/lib/schema.py

@@ -1,4 +1,5 @@
 """ schema format representation """
+import abc
 import typing
 from dataclasses import dataclass, field
 from typing import List, Set, Union, Dict, Optional
@@ -136,40 +137,29 @@ predicate_marker = object()
 TypeRef = Union[type, str]
 
 
-@functools.singledispatch
 def get_type_name(arg: TypeRef) -> str:
-    raise Error(f"Not a schema type or string ({arg})")
+    match arg:
+        case type():
+            return arg.__name__
+        case str():
+            return arg
+        case _:
+            raise Error(f"Not a schema type or string ({arg})")
 
 
-@get_type_name.register
-def _(arg: type):
-    return arg.__name__
-
-
-@get_type_name.register
-def _(arg: str):
-    return arg
-
-
-@functools.singledispatch
 def _make_property(arg: object) -> Property:
-    if arg is predicate_marker:
-        return PredicateProperty()
-    raise Error(f"Illegal property specifier {arg}")
+    match arg:
+        case _ if arg is predicate_marker:
+            return PredicateProperty()
+        case str() | type():
+            return SingleProperty(type=get_type_name(arg))
+        case Property():
+            return arg
+        case _:
+            raise Error(f"Illegal property specifier {arg}")
 
 
-@_make_property.register(str)
-@_make_property.register(type)
-def _(arg: TypeRef):
-    return SingleProperty(type=get_type_name(arg))
-
-
-@_make_property.register
-def _(arg: Property):
-    return arg
-
-
-class PropertyModifier:
+class PropertyModifier(abc.ABC):
     """ Modifier of `Property` objects.
         Being on the right of `|` it will trigger construction of a `Property` from
         the left operand.
@@ -180,8 +170,14 @@ class PropertyModifier:
         self.modify(ret)
         return ret
 
+    def __invert__(self) -> "PropertyModifier":
+        return self.negate()
+
     def modify(self, prop: Property):
-        raise NotImplementedError
+        ...
+
+    def negate(self) -> "PropertyModifier":
+        ...
 
 
 def split_doc(doc):

misc/codegen/lib/schemadefs.py

@@ -1,40 +1,66 @@
-from typing import Callable as _Callable
+from typing import Callable as _Callable, List as _List
 from misc.codegen.lib import schema as _schema
 import inspect as _inspect
 from dataclasses import dataclass as _dataclass
 
+from misc.codegen.lib.schema import Property
 
+
+@_dataclass
 class _ChildModifier(_schema.PropertyModifier):
+    child: bool = True
+
     def modify(self, prop: _schema.Property):
         if prop.type is None or prop.type[0].islower():
             raise _schema.Error("Non-class properties cannot be children")
         if prop.is_unordered:
             raise _schema.Error("Set properties cannot be children")
-        prop.is_child = True
+        prop.is_child = self.child
+
+    def negate(self) -> _schema.PropertyModifier:
+        return _ChildModifier(False)
+
+# make ~doc same as doc(None)
+
+
+class _DocModifierMetaclass(type(_schema.PropertyModifier)):
+    def __invert__(self) -> _schema.PropertyModifier:
+        return _DocModifier(None)
 
 
 @_dataclass
-class _DocModifier(_schema.PropertyModifier):
-    doc: str
+class _DocModifier(_schema.PropertyModifier, metaclass=_DocModifierMetaclass):
+    doc: str | None
 
     def modify(self, prop: _schema.Property):
-        if "\n" in self.doc or self.doc[-1] == ".":
+        if self.doc and ("\n" in self.doc or self.doc[-1] == "."):
             raise _schema.Error("No newlines or trailing dots are allowed in doc, did you intend to use desc?")
         prop.doc = self.doc
 
+    def negate(self) -> _schema.PropertyModifier:
+        return _DocModifier(None)
+
+
+# make ~desc same as desc(None)
+class _DescModifierMetaclass(type(_schema.PropertyModifier)):
+    def __invert__(self) -> _schema.PropertyModifier:
+        return _DescModifier(None)
+
 
 @_dataclass
-class _DescModifier(_schema.PropertyModifier):
-    description: str
+class _DescModifier(_schema.PropertyModifier, metaclass=_DescModifierMetaclass):
+    description: str | None
 
     def modify(self, prop: _schema.Property):
         prop.description = _schema.split_doc(self.description)
 
+    def negate(self) -> _schema.PropertyModifier:
+        return _DescModifier(None)
+
 
 def include(source: str):
     # add to `includes` variable in calling context
-    _inspect.currentframe().f_back.f_locals.setdefault(
-        "__includes", []).append(source)
+    _inspect.currentframe().f_back.f_locals.setdefault("includes", []).append(source)
 
 
 class _Namespace:
@@ -44,9 +70,15 @@ class _Namespace:
         self.__dict__.update(kwargs)
 
 
+@_dataclass
 class _SynthModifier(_schema.PropertyModifier, _Namespace):
+    synth: bool = True
+
     def modify(self, prop: _schema.Property):
-        prop.synth = True
+        prop.synth = self.synth
+
+    def negate(self) -> "PropertyModifier":
+        return _SynthModifier(False)
 
 
 qltest = _Namespace()
@@ -63,22 +95,35 @@ class _Pragma(_schema.PropertyModifier):
     For schema classes it acts as a python decorator with `@`.
     """
     pragma: str
+    remove: bool = False
 
     def __post_init__(self):
         namespace, _, name = self.pragma.partition('_')
         setattr(globals()[namespace], name, self)
 
     def modify(self, prop: _schema.Property):
-        prop.pragmas.append(self.pragma)
+        self._apply(prop.pragmas)
+
+    def negate(self) -> "PropertyModifier":
+        return _Pragma(self.pragma, remove=True)
 
     def __call__(self, cls: type) -> type:
         """ use this pragma as a decorator on classes """
         if "_pragmas" in cls.__dict__:  # not using hasattr as we don't want to land on inherited pragmas
-            cls._pragmas.append(self.pragma)
-        else:
+            self._apply(cls._pragmas)
+        elif not self.remove:
             cls._pragmas = [self.pragma]
         return cls
 
+    def _apply(self, pragmas: _List[str]) -> None:
+        if self.remove:
+            try:
+                pragmas.remove(self.pragma)
+            except ValueError:
+                pass
+        else:
+            pragmas.append(self.pragma)
+
 
 class _Optionalizer(_schema.PropertyModifier):
     def modify(self, prop: _schema.Property):
@@ -172,7 +217,28 @@ synth.on_arguments = lambda **kwargs: _annotate(
     synth=_schema.SynthInfo(on_arguments={k: _schema.get_type_name(t) for k, t in kwargs.items()}))
 
 
-def annotate(annotated_cls: type) -> _Callable[[type], None]:
+class _PropertyModifierList(_schema.PropertyModifier):
+    def __init__(self):
+        self._mods = []
+
+    def __or__(self, other: _schema.PropertyModifier):
+        self._mods.append(other)
+        return self
+
+    def modify(self, prop: Property):
+        for m in self._mods:
+            m.modify(prop)
+
+
+class _PropertyAnnotation:
+    def __or__(self, other: _schema.PropertyModifier):
+        return _PropertyModifierList() | other
+
+
+_ = _PropertyAnnotation()
+
+
+def annotate(annotated_cls: type) -> _Callable[[type], _PropertyAnnotation]:
     """
     Add or modify schema annotations after a class has been defined
     For the moment, only docstring annotation is supported. In the future, any kind of
@@ -180,9 +246,17 @@ def annotate(annotated_cls: type) -> _Callable[[type], None]:
 
     The name of the class used for annotation must be `_`
     """
-    def decorator(cls: type) -> None:
+    def decorator(cls: type) -> _PropertyAnnotation:
         if cls.__name__ != "_":
             raise _schema.Error("Annotation classes must be named _")
         annotated_cls.__doc__ = cls.__doc__
-        return None
+        for p, a in cls.__annotations__.items():
+            if p in annotated_cls.__annotations__:
+                annotated_cls.__annotations__[p] |= a
+            elif isinstance(a, (_PropertyAnnotation, _PropertyModifierList)):
+                raise _schema.Error(f"annotated property {p} not present in annotated class "
+                                    f"{annotated_cls.__name__}")
+            else:
+                annotated_cls.__annotations__[p] = a
+        return _
     return decorator

misc/codegen/loaders/schemaloader.py

@@ -136,7 +136,7 @@ def load(m: types.ModuleType) -> schema.Schema:
     for name, data in m.__dict__.items():
         if hasattr(defs, name):
             continue
-        if name == "__includes":
+        if name == "includes":
             includes = data
             continue
         if name.startswith("__") or name == "_":

misc/codegen/test/test_schemaloader.py

@@ -1,6 +1,7 @@
 import sys
 
 import pytest
+from misc.codegen.lib.schemadefs import optional
 
 from misc.codegen.test.utils import *
 from misc.codegen.lib import schemadefs as defs
@@ -774,6 +775,63 @@ def test_annotate_docstring():
     }
 
 
+def test_annotate_fields():
+    @load
+    class data:
+        class Root:
+            x: defs.int
+            y: defs.optional["Root"] | defs.child
+
+        @defs.annotate(Root)
+        class _:
+            x: defs._ | defs.doc("foo")
+            y: defs._ | defs.ql.internal
+            z: defs.string
+
+    assert data.classes == {
+        "Root": schema.Class("Root", properties=[
+            schema.SingleProperty("x", "int", doc="foo"),
+            schema.OptionalProperty("y", "Root", pragmas=["ql_internal"], is_child=True),
+            schema.SingleProperty("z", "string"),
+        ]),
+    }
+
+
+def test_annotate_fields_negations():
+    @load
+    class data:
+        class Root:
+            x: defs.int | defs.ql.internal | defs.qltest.skip
+            y: defs.optional["Root"] | defs.child | defs.desc("foo\nbar\n")
+            z: defs.string | defs.synth | defs.doc("foo")
+
+        @defs.annotate(Root)
+        class _:
+            x: defs._ | ~defs.ql.internal
+            y: defs._ | ~defs.child | ~defs.ql.internal | ~defs.desc
+            z: defs._ | ~defs.synth | ~defs.doc
+
+    assert data.classes == {
+        "Root": schema.Class("Root", properties=[
+            schema.SingleProperty("x", "int", pragmas=["qltest_skip"]),
+            schema.OptionalProperty("y", "Root"),
+            schema.SingleProperty("z", "string"),
+        ]),
+    }
+
+
+def test_annotate_non_existing_field():
+    with pytest.raises(schema.Error):
+        @load
+        class data:
+            class Root:
+                pass
+
+            @defs.annotate(Root)
+            class _:
+                x: defs._ | defs.doc("foo")
+
+
 def test_annotate_not_underscore():
     with pytest.raises(schema.Error):
         @load

rust/schema/__init__.py

@@ -11,6 +11,3 @@ For how documentation of generated QL code works, please read `misc/codegen/sche
 
 from .prelude import *
 from .ast import *
-
-include("../shared/tree-sitter-extractor/src/generator/prefix.dbscheme")
-include("prefix.dbscheme")

rust/schema/prelude.py

@@ -1,5 +1,8 @@
 from misc.codegen.lib.schemadefs import *
 
+include("../shared/tree-sitter-extractor/src/generator/prefix.dbscheme")
+include("prefix.dbscheme")
+
 @qltest.skip
 class Element:
     pass