Merge pull request #17530 from github/redsun82/codegen-annotate

Codegen: allow full annotation of classes

Author: aibaars

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 @@ def null_class(self):
 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}")
-
-
-@_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:
+    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}")
+
+
+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 @@ def __ror__(self, other: object) -> Property:
         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,65 @@
-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)
+
+
+class _DocModifierMetaclass(type(_schema.PropertyModifier)):
+    # make ~doc same as doc(None)
+    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)
+
+
+class _DescModifierMetaclass(type(_schema.PropertyModifier)):
+    # make ~desc same as desc(None)
+    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 +69,15 @@ def __init__(self, **kwargs):
         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 +94,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,17 +216,57 @@ def group(name: str = "") -> _ClassDecorator:
     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
     modification will be allowed.
 
     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
+        if cls.__doc__ is not None:
+            annotated_cls.__doc__ = cls.__doc__
+        old_pragmas = getattr(annotated_cls, "_pragmas", None)
+        new_pragmas = getattr(cls, "_pragmas", [])
+        if old_pragmas:
+            old_pragmas.extend(new_pragmas)
+        else:
+            annotated_cls._pragmas = new_pragmas
+        for a, v in cls.__dict__.items():
+            # transfer annotations
+            if a.startswith("_") and not a.startswith("__") and a != "_pragmas":
+                setattr(annotated_cls, a, v)
+        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
@@ -762,18 +763,103 @@ class data:
         class Root:
             """ old docstring """
 
+        class A(Root):
+            """ A docstring """
+
         @defs.annotate(Root)
         class _:
             """
             new
             docstring
             """
 
+        @defs.annotate(A)
+        class _:
+            pass
+
+    assert data.classes == {
+        "Root": schema.Class("Root", doc=["new", "docstring"], derived={"A"}),
+        "A": schema.Class("A", bases=["Root"], doc=["A docstring"]),
+    }
+
+
+def test_annotate_decorations():
+    @load
+    class data:
+        @defs.qltest.skip
+        class Root:
+            pass
+
+        @defs.annotate(Root)
+        @defs.qltest.collapse_hierarchy
+        @defs.ql.hideable
+        @defs.cpp.skip
+        class _:
+            pass
+
+    assert data.classes == {
+        "Root": schema.Class("Root", hideable=True,
+                             pragmas=["qltest_skip", "cpp_skip", "qltest_collapse_hierarchy"]),
+    }
+
+
+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", doc=["new", "docstring"]),
+        "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