Tool for generating schema type annotations

edb/schema/annos.py

@@ -43,12 +43,14 @@
 from . import referencing
 from . import objects as so
 from . import utils
+from .generated import annos as sg_annos
 
 if TYPE_CHECKING:
     from . import schema as s_schema
 
 
 class AnnotationValue(
+    sg_annos.AnnotationValueMixin,
     referencing.ReferencedInheritingObject,
     qlkind=qltypes.SchemaObjectClass.ANNOTATION,
     reflection=so.ReflectionMethod.AS_LINK,
@@ -101,7 +103,7 @@ def get_verbosename(
 T = TypeVar("T")
 
 
-class AnnotationSubject(so.Object):
+class AnnotationSubject(so.Object, sg_annos.AnnotationSubjectMixin):
 
     annotations_refs = so.RefDict(
         attr='annotations',
@@ -177,6 +179,7 @@ def must_get_json_annotation(
 
 
 class Annotation(
+    sg_annos.AnnotationMixin,
     so.QualifiedObject,
     so.InheritingObject,
     AnnotationSubject,

edb/schema/casts.py

@@ -36,6 +36,7 @@
 from . import types as s_types
 from . import schema as s_schema
 from . import utils
+from .generated import casts as sg_casts
 
 
 _NOT_REACHABLE = 10000000
@@ -206,6 +207,7 @@ def get_cast_fullname(
 
 
 class Cast(
+    sg_casts.CastMixin,
     so.QualifiedObject,
     s_anno.AnnotationSubject,
     s_func.VolatilitySubject,

edb/schema/constraints.py

@@ -57,6 +57,7 @@
 from . import pseudo as s_pseudo
 from . import referencing
 from . import utils
+from .generated import constraints as sg_constraints
 
 
 if TYPE_CHECKING:
@@ -122,8 +123,10 @@ def get_key_for_name(
 
 
 class Constraint(
+    sg_constraints.ConstraintMixin,
     referencing.ReferencedInheritingObject,
-    s_func.CallableObject, s_abc.Constraint,
+    s_func.CallableObject,
+    s_abc.Constraint,
     qlkind=ft.SchemaObjectClass.CONSTRAINT,
     data_safe=True,
 ):
@@ -364,6 +367,7 @@ def get_default_base_name(self) -> sn.QualName:
 
 
 class ConsistencySubject(
+    sg_constraints.ConsistencySubjectMixin,
     so.QualifiedObject,
     so.InheritingObject,
     s_anno.AnnotationSubject,

edb/schema/database.py

@@ -32,11 +32,13 @@
 from . import delta as sd
 from . import objects as so
 from . import schema as s_schema
+from .generated import database as sg_database
 
 from typing import cast
 
 
 class Database(
+    sg_database.DatabaseMixin,
     so.ExternalObject,
     s_anno.AnnotationSubject,
     s_abc.Database,

edb/schema/expraliases.py

@@ -34,6 +34,7 @@
 from . import name as sn
 from . import objects as so
 from . import types as s_types
+from .generated import expraliases as sg_expraliases
 
 
 if TYPE_CHECKING:
@@ -42,6 +43,7 @@
 
 
 class Alias(
+    sg_expraliases.AliasMixin,
     so.QualifiedObject,
     s_anno.AnnotationSubject,
     qlkind=qltypes.SchemaObjectClass.ALIAS,

edb/schema/extensions.py

@@ -39,9 +39,11 @@
 from . import name as sn
 from . import objects as so
 from . import schema as s_schema
+from .generated import extensions as sg_extensions
 
 
 class ExtensionPackage(
+    sg_extensions.ExtensionPackageMixin,
     so.GlobalObject,
     s_anno.AnnotationSubject,
     qlkind=qltypes.SchemaObjectClass.EXTENSION_PACKAGE,
@@ -102,6 +104,7 @@ def get_displayname_static(cls, name: sn.Name) -> str:
 
 
 class Extension(
+    sg_extensions.ExtensionMixin,
     so.Object,
     qlkind=qltypes.SchemaObjectClass.EXTENSION,
     data_safe=False,

edb/schema/functions.py

@@ -62,6 +62,7 @@
 from . import referencing
 from . import types as s_types
 from . import utils
+from .generated import functions as sg_functions
 
 
 if TYPE_CHECKING:
@@ -151,32 +152,8 @@ def param_is_inherited(
     return qualname != param_name.name
 
 
-class ParameterLike(s_abc.Parameter):
-
-    def get_parameter_name(self, schema: s_schema.Schema) -> str:
-        raise NotImplementedError
-
-    def get_name(self, schema: s_schema.Schema) -> sn.Name:
-        raise NotImplementedError
-
-    def get_kind(self, _: s_schema.Schema) -> ft.ParameterKind:
-        raise NotImplementedError
-
-    def get_default(self, _: s_schema.Schema) -> Optional[s_expr.Expression]:
-        raise NotImplementedError
-
-    def get_type(self, _: s_schema.Schema) -> s_types.Type:
-        raise NotImplementedError
-
-    def get_typemod(self, _: s_schema.Schema) -> ft.TypeModifier:
-        raise NotImplementedError
-
-    def as_str(self, schema: s_schema.Schema) -> str:
-        raise NotImplementedError
-
-
 # Non-schema description of a parameter.
-class ParameterDesc(ParameterLike):
+class ParameterDesc(s_abc.Parameter):
 
     num: int
     name: sn.Name
@@ -364,9 +341,10 @@ def make_func_param(
 
 
 class Parameter(
+    sg_functions.ParameterMixin,
     so.ObjectFragment,
     so.Object,  # Help reflection figure out the right db MRO
-    ParameterLike,
+    s_abc.Parameter,
     qlkind=ft.SchemaObjectClass.PARAMETER,
     data_safe=True,
 ):
@@ -489,6 +467,9 @@ def get_ast(self, schema: s_schema.Schema) -> qlast.FuncParam:
         )
 
 
+ParameterLike = ParameterDesc | Parameter
+
+
 class CallableCommandContext(sd.ObjectCommandContext['CallableObject'],
                              s_anno.AnnotationSubjectCommandContext):
     pass
@@ -756,7 +737,7 @@ def compare_values(
         return 1.0
 
 
-class VolatilitySubject(so.Object):
+class VolatilitySubject(so.Object, sg_functions.VolatilitySubjectMixin):
 
     volatility = so.SchemaField(
         ft.Volatility, default=ft.Volatility.Volatile,
@@ -792,6 +773,7 @@ def get_abstract(self, schema: s_schema.Schema) -> bool:
 
 
 class CallableObject(
+    sg_functions.CallableObjectMixin,
     so.QualifiedObject,
     s_anno.AnnotationSubject,
     CallableLike,
@@ -1226,6 +1208,7 @@ def _delete_begin(
 
 
 class Function(
+    sg_functions.FunctionMixin,
     CallableObject,
     VolatilitySubject,
     s_abc.Function,

edb/schema/futures.py

@@ -29,9 +29,11 @@
 from . import objects as so
 from . import name as sn
 from . import schema as s_schema
+from .generated import futures as sg_futures
 
 
 class FutureBehavior(
+    sg_futures.FutureBehaviorMixin,
     so.Object,
     qlkind=qltypes.SchemaObjectClass.FUTURE,
     data_safe=False,

edb/schema/generated/annos.py

@@ -0,0 +1,57 @@
+# DO NOT EDIT. This file was generated with:
+#
+# $ edb gen-schema-mixins
+
+"""Type definitions for generated methods on schema classes"""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edb.schema import schema as s_schema
+from edb.schema import orm as s_orm
+from edb.schema import objects
+from edb.schema import annos
+
+
+class AnnotationValueMixin:
+
+    def get_subject(
+        self, schema: 's_schema.Schema'
+    ) -> 'objects.Object':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'subject'    # type: ignore
+        )
+
+    def get_annotation(
+        self, schema: 's_schema.Schema'
+    ) -> 'annos.Annotation':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'annotation'    # type: ignore
+        )
+
+    def get_value(
+        self, schema: 's_schema.Schema'
+    ) -> 'str':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'value'    # type: ignore
+        )
+
+
+class AnnotationSubjectMixin:
+
+    def get_annotations(
+        self, schema: 's_schema.Schema'
+    ) -> 'objects.ObjectIndexByShortname[annos.AnnotationValue]':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'annotations'    # type: ignore
+        )
+
+
+class AnnotationMixin:
+
+    def get_inheritable(
+        self, schema: 's_schema.Schema'
+    ) -> 'bool':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'inheritable'    # type: ignore
+        )

edb/schema/generated/casts.py

@@ -0,0 +1,79 @@
+# DO NOT EDIT. This file was generated with:
+#
+# $ edb gen-schema-mixins
+
+"""Type definitions for generated methods on schema classes"""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edb.schema import schema as s_schema
+from edb.schema import orm as s_orm
+from edb.edgeql import ast
+from edb.schema import types
+
+
+class CastMixin:
+
+    def get_from_type(
+        self, schema: 's_schema.Schema'
+    ) -> 'types.Type':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'from_type'    # type: ignore
+        )
+
+    def get_to_type(
+        self, schema: 's_schema.Schema'
+    ) -> 'types.Type':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'to_type'    # type: ignore
+        )
+
+    def get_allow_implicit(
+        self, schema: 's_schema.Schema'
+    ) -> 'bool':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'allow_implicit'    # type: ignore
+        )
+
+    def get_allow_assignment(
+        self, schema: 's_schema.Schema'
+    ) -> 'bool':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'allow_assignment'    # type: ignore
+        )
+
+    def get_language(
+        self, schema: 's_schema.Schema'
+    ) -> 'ast.Language':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'language'    # type: ignore
+        )
+
+    def get_from_function(
+        self, schema: 's_schema.Schema'
+    ) -> 'str':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'from_function'    # type: ignore
+        )
+
+    def get_from_expr(
+        self, schema: 's_schema.Schema'
+    ) -> 'bool':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'from_expr'    # type: ignore
+        )
+
+    def get_from_cast(
+        self, schema: 's_schema.Schema'
+    ) -> 'bool':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'from_cast'    # type: ignore
+        )
+
+    def get_code(
+        self, schema: 's_schema.Schema'
+    ) -> 'str':
+        return s_orm.get_field_value(  # type: ignore
+            self, schema, 'code'    # type: ignore
+        )