codeql/misc/codegen/lib/schemadefs.py

from typing import Callable as _Callable, Dict as _Dict, ClassVar as _ClassVar
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

_set = set


@_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 = 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, metaclass=_DocModifierMetaclass):
    doc: str | None

    def modify(self, prop: _schema.Property):
        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, 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)


@_dataclass
class _Namespace:
    """ simple namespacing mechanism """
    name: str

    def add(self, pragma: "_PragmaBase", key: str | None = None):
        self.__dict__[pragma.pragma] = pragma
        pragma.pragma = key or f"{self.name}_{pragma.pragma}"


@_dataclass
class _SynthModifier(_schema.PropertyModifier, _Namespace):
    synth: bool = True

    def modify(self, prop: _schema.Property):
        prop.synth = self.synth

    def negate(self) -> _schema.PropertyModifier:
        return _SynthModifier(self.name, False)


qltest = _Namespace("qltest")
ql = _Namespace("ql")
cpp = _Namespace("cpp")
rust = _Namespace("rust")
synth = _SynthModifier("synth")


@_dataclass
class _PragmaBase:
    pragma: str


@_dataclass
class _ClassPragma(_PragmaBase):
    """ A class pragma.
    For schema classes it acts as a python decorator with `@`.
    """
    inherited: bool = False
    value: object = None

    def __call__(self, cls: type) -> type:
        """ use this pragma as a decorator on classes """
        if self.inherited:
            setattr(cls, f"{_schema.inheritable_pragma_prefix}{self.pragma}", self.value)
        else:
            # not using hasattr as we don't want to land on inherited pragmas
            if "_pragmas" not in cls.__dict__:
                cls._pragmas = {}
            self._apply(cls._pragmas)
        return cls

    def _apply(self, pragmas: _Dict[str, object]) -> None:
        pragmas[self.pragma] = self.value


@_dataclass
class _Pragma(_ClassPragma, _schema.PropertyModifier):
    """ A class or property pragma.
    For properties, it functions similarly to a `_PropertyModifier` with `|`, adding the pragma.
    For schema classes it acts as a python decorator with `@`.
    """
    remove: bool = False

    def modify(self, prop: _schema.Property):
        self._apply(prop.pragmas)

    def negate(self) -> _schema.PropertyModifier:
        return _Pragma(self.pragma, remove=True)

    def _apply(self, pragmas: _Dict[str, object]) -> None:
        if self.remove:
            pragmas.pop(self.pragma, None)
        else:
            super()._apply(pragmas)


@_dataclass
class _ParametrizedClassPragma(_PragmaBase):
    """ A class parametrized pragma.
    Needs to be applied to a parameter to give a class pragma.
    """
    _pragma_class: _ClassVar[type] = _ClassPragma

    inherited: bool = False
    factory: _Callable[..., object] = None

    def __post_init__(self):
        self.__signature__ = _inspect.signature(self.factory).replace(return_annotation=self._pragma_class)

    def __call__(self, *args, **kwargs) -> _pragma_class:
        return self._pragma_class(self.pragma, self.inherited, value=self.factory(*args, **kwargs))


@_dataclass
class _ParametrizedPragma(_ParametrizedClassPragma):
    """ A class or property parametrized pragma.
    Needs to be applied to a parameter to give a pragma.
    """
    _pragma_class: _ClassVar[type] = _Pragma

    def __invert__(self) -> _Pragma:
        return _Pragma(self.pragma, remove=True)


class _Optionalizer(_schema.PropertyModifier):
    def modify(self, prop: _schema.Property):
        K = _schema.Property.Kind
        if prop.kind != K.SINGLE:
            raise _schema.Error(
                "optional should only be applied to simple property types")
        prop.kind = K.OPTIONAL


class _Listifier(_schema.PropertyModifier):
    def modify(self, prop: _schema.Property):
        K = _schema.Property.Kind
        if prop.kind == K.SINGLE:
            prop.kind = K.REPEATED
        elif prop.kind == K.OPTIONAL:
            prop.kind = K.REPEATED_OPTIONAL
        else:
            raise _schema.Error(
                "list should only be applied to simple or optional property types")


class _Setifier(_schema.PropertyModifier):
    def modify(self, prop: _schema.Property):
        K = _schema.Property.Kind
        if prop.kind != K.SINGLE:
            raise _schema.Error("set should only be applied to simple property types")
        prop.kind = K.REPEATED_UNORDERED


class _TypeModifier:
    """ Modifies types using get item notation """

    def __init__(self, modifier: _schema.PropertyModifier):
        self.modifier = modifier

    def __getitem__(self, item):
        return item | self.modifier


_ClassDecorator = _Callable[[type], type]


def _annotate(**kwargs) -> _ClassDecorator:
    def f(cls: type) -> type:
        for k, v in kwargs.items():
            setattr(cls, f"_{k}", v)
        return cls

    return f


boolean = "boolean"
int = "int"
string = "string"

predicate = _schema.predicate_marker
optional = _TypeModifier(_Optionalizer())
list = _TypeModifier(_Listifier())
set = _TypeModifier(_Setifier())

child = _ChildModifier()
doc = _DocModifier
desc = _DescModifier

use_for_null = _annotate(null=True)

qltest.add(_Pragma("skip"))
qltest.add(_ClassPragma("collapse_hierarchy"))
qltest.add(_ClassPragma("uncollapse_hierarchy"))
qltest.add(_ParametrizedClassPragma("test_with", inherited=True, factory=_schema.get_type_name))

ql.add(_ParametrizedClassPragma("default_doc_name", factory=lambda doc: doc))
ql.add(_ClassPragma("hideable", inherited=True))
ql.add(_Pragma("internal"))

cpp.add(_Pragma("skip"))

rust.add(_Pragma("skip_doc_test"))

rust.add(_ParametrizedClassPragma("doc_test_signature", factory=lambda signature: signature))


def group(name: str = "") -> _ClassDecorator:
    return _annotate(group=name)


synth.add(_ParametrizedClassPragma("from_class", factory=lambda ref: _schema.SynthInfo(
    from_class=_schema.get_type_name(ref))), key="synth")
synth.add(_ParametrizedClassPragma("on_arguments", factory=lambda **kwargs:
                                   _schema.SynthInfo(on_arguments={k: _schema.get_type_name(t) for k, t in kwargs.items()})), key="synth")


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) -> _PropertyAnnotation:
        if cls.__name__ != "_":
            raise _schema.Error("Annotation classes must be named _")
        if cls.__doc__ is not None:
            annotated_cls.__doc__ = cls.__doc__
        for p, v in cls.__dict__.get("_pragmas", {}).items():
            _ClassPragma(p, value=v)(annotated_cls)
        for a in dir(cls):
            if a.startswith(_schema.inheritable_pragma_prefix):
                setattr(annotated_cls, a, getattr(cls, a))
        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