diff --git a/misc/codegen/lib/schema.py b/misc/codegen/lib/schema.py index 025c1330cce..b950298fd97 100644 --- a/misc/codegen/lib/schema.py +++ b/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): diff --git a/misc/codegen/lib/schemadefs.py b/misc/codegen/lib/schemadefs.py index ed899e71171..d5d7576c5b4 100644 --- a/misc/codegen/lib/schemadefs.py +++ b/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 diff --git a/misc/codegen/loaders/schemaloader.py b/misc/codegen/loaders/schemaloader.py index 7720c4057d0..d00eb2b7ce7 100644 --- a/misc/codegen/loaders/schemaloader.py +++ b/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 == "_": diff --git a/misc/codegen/test/test_schemaloader.py b/misc/codegen/test/test_schemaloader.py index 1880454a8e7..02d8f71d41c 100644 --- a/misc/codegen/test/test_schemaloader.py +++ b/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 diff --git a/rust/schema/__init__.py b/rust/schema/__init__.py index 4bfc8715d51..552a55f9480 100644 --- a/rust/schema/__init__.py +++ b/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") diff --git a/rust/schema/prelude.py b/rust/schema/prelude.py index 2272af01778..df4cd1dac7e 100644 --- a/rust/schema/prelude.py +++ b/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