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

Codegen: allow to attach docstrings after the definition

misc/codegen/lib/schemadefs.py

@@ -177,3 +177,19 @@ synth.from_class = lambda ref: _annotate(synth=_schema.SynthInfo(
     from_class=_schema.get_type_name(ref)))
 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]:
+    """
+    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:
+        if cls.__name__ != "_":
+            raise _schema.Error("Annotation classes must be named _")
+        annotated_cls.__doc__ = cls.__doc__
+        return None
+    return decorator

misc/codegen/loaders/schemaloader.py

@@ -138,7 +138,7 @@ def load(m: types.ModuleType) -> schema.Schema:
         if name == "__includes":
             includes = data
             continue
-        if name.startswith("__"):
+        if name.startswith("__") or name == "_":
             continue
         cls = _get_class(data)
         if classes and not cls.bases:

misc/codegen/test/test_schemaloader.py

@@ -756,6 +756,39 @@ def test_test_with():
     }
 
 
+def test_annotate_docstring():
+    @load
+    class data:
+        class Root:
+            """ old docstring """
+
+        @defs.annotate(Root)
+        class _:
+            """
+            new
+            docstring
+            """
+
+    assert data.classes == {
+        "Root": schema.Class("Root", doc=["new", "docstring"]),
+    }
+
+
+def test_annotate_not_underscore():
+    with pytest.raises(schema.Error):
+        @load
+        class data:
+            class Root:
+                pass
+
+            @defs.annotate(Root)
+            class Something:
+                """
+                new
+                docstring
+                """
+
+
 def test_test_with_unknown_string():
     with pytest.raises(schema.Error):
         @load