Improve generated type hints and defaults (#606)

* Update codegen for better type hints and default

* No longer needed to ignor lint rule

* codegen and adjust

* Add missing struct checks

* Some args must be optional, because the descriptor is, or because we want them to

* Comments

* fragment should be optional too

codegen/apipatcher.py

@@ -4,10 +4,32 @@
 """
 
 from codegen.utils import print, format_code, to_snake_case, to_camel_case, Patcher
-from codegen.idlparser import get_idl_parser
+from codegen.idlparser import get_idl_parser, Attribute
 from codegen.files import file_cache
 
 
+# In wgpu-py, we make some args optional, that are not optional in the
+# IDL. Reasons may be because it makes sense to be able to omit them,
+# or because the WebGPU says its optional while IDL says its not, or
+# for backwards compatibility. These args have a default value of
+# 'optional'  (which is just None) so we can recognise them. If IDL
+# makes any of these args optional, their presense in this list is
+# ignored.
+ARGS_TO_MAKE_OPTIONAL = {
+    ("compilation_hints", "compilation_hints"),  # idl actually has a default
+    ("create_shader_module", "source_map"),
+    ("begin_compute_pass", "timestamp_writes"),
+    ("begin_render_pass", "timestamp_writes"),
+    ("begin_render_pass", "depth_stencil_attachment"),
+    ("begin_render_pass", "occlusion_query_set"),
+    ("create_render_pipeline", "depth_stencil"),
+    ("create_render_pipeline", "fragment"),
+    ("create_render_pipeline_async", "depth_stencil"),
+    ("create_render_pipeline_async", "fragment"),
+    ("create_render_bundle_encoder", "depth_stencil_format"),
+}
+
+
 def patch_base_api(code):
     """Given the Python code, applies patches to make the code conform
     to the IDL.
@@ -406,49 +428,57 @@ def get_method_def(self, classname, methodname):
         # Get arg names and types
         idl_line = functions[name_idl]
         args = idl_line.split("(", 1)[1].split(")", 1)[0].split(",")
-        args = [arg.strip() for arg in args if arg.strip()]
-        raw_defaults = [arg.partition("=")[2].strip() for arg in args]
-        place_holder_default = False
-        defaults = []
-        for default, arg in zip(raw_defaults, args):
-            if default:
-                place_holder_default = "None"  # any next args must have a default
-            elif arg.startswith("optional "):
-                default = "None"
-            else:
-                default = place_holder_default
-            defaults.append(default)
-
-        argnames = [arg.split("=")[0].split()[-1] for arg in args]
-        argnames = [to_snake_case(argname) for argname in argnames]
-        argnames = [(f"{n}={v}" if v else n) for n, v in zip(argnames, defaults)]
-        argtypes = [arg.split("=")[0].split()[-2] for arg in args]
+        args = [Attribute(arg) for arg in args if arg.strip()]
 
         # If one arg that is a dict, flatten dict to kwargs
-        if len(argtypes) == 1 and argtypes[0].endswith(
+        if len(args) == 1 and args[0].typename.endswith(
             ("Options", "Descriptor", "Configuration")
         ):
-            assert argtypes[0].startswith("GPU")
-            fields = self.idl.structs[argtypes[0][3:]].values()  # struct fields
-            py_args = [self._arg_from_struct_field(field) for field in fields]
+            assert args[0].typename.startswith("GPU")
+            des_is_optional = bool(args[0].default)
+            attributes = self.idl.structs[args[0].typename[3:]].values()
+            py_args = [
+                self._arg_from_attribute(methodname, attr, des_is_optional)
+                for attr in attributes
+            ]
             if py_args[0].startswith("label: str"):
-                py_args[0] = 'label=""'
+                py_args[0] = 'label: str=""'
             py_args = ["self", "*", *py_args]
         else:
-            py_args = ["self", *argnames]
+            py_args = [self._arg_from_attribute(methodname, attr) for attr in args]
+            py_args = ["self", *py_args]
+
+            # IDL has some signatures that cannot work in Python. This may be a bug in idl
+            known_bugged_methods = {"GPUPipelineError.__init__"}
+            remove_default = False
+            for i in reversed(range(len(py_args))):
+                arg = py_args[i]
+                if "=" in arg:
+                    if remove_default:
+                        py_args[i] = arg.split("=")[0]
+                        assert f"{classname}.{methodname}" in known_bugged_methods
+                else:
+                    remove_default = True
 
         # Construct final def
         line = preamble + ", ".join(py_args) + "): pass\n"
         line = format_code(line, True).split("):")[0] + "):"
         return "    " + line
 
-    def _arg_from_struct_field(self, field):
-        name = to_snake_case(field.name)
-        d = field.default
-        t = self.idl.resolve_type(field.typename)
+    def _arg_from_attribute(self, methodname, attribute, force_optional=False):
+        name = to_snake_case(attribute.name)
+        optional_in_py = (methodname, name) in ARGS_TO_MAKE_OPTIONAL
+        d = attribute.default
+        t = self.idl.resolve_type(attribute.typename)
         result = name
+        if (force_optional or optional_in_py) and not d:
+            d = "optional"
         if t:
-            result += f": {t}"
+            # If default is None, the type won't match, so we need to mark it optional
+            if d == "None":
+                result += f": Optional[{t}]"
+            else:
+                result += f": {t}"
         if d:
             d = {"false": "False", "true": "True"}.get(d, d)
             result += f"={d}"
@@ -504,7 +534,20 @@ def get_method_comment(self, classname, methodname):
         functions = self.idl.classes[classname].functions
         name_idl = self.method_is_known(classname, methodname)
         if name_idl:
-            return "    # IDL: " + functions[name_idl]
+            idl_line = functions[name_idl]
+
+            args = idl_line.split("(", 1)[1].split(")", 1)[0].split(",")
+            args = [Attribute(arg) for arg in args if arg.strip()]
+
+            # If one arg that is a dict, flatten dict to kwargs
+            if len(args) == 1 and args[0].typename.endswith(
+                ("Options", "Descriptor", "Configuration")
+            ):
+                assert args[0].typename.startswith("GPU")
+                attributes = self.idl.structs[args[0].typename[3:]].values()
+                idl_line += " -> " + ", ".join(attr.line for attr in attributes)
+
+            return "    # IDL: " + idl_line
 
 
 class BackendApiPatcher(AbstractApiPatcher):
@@ -575,7 +618,7 @@ def apply(self, code):
 
         idl = get_idl_parser()
         all_structs = set()
-        ignore_structs = {"Extent3D"}
+        ignore_structs = {"Extent3D", "Origin3D"}
 
         for classname, i1, i2 in self.iter_classes():
             if classname not in idl.classes:
@@ -621,8 +664,8 @@ def apply(self, code):
 
     def _get_sub_structs(self, idl, structname):
         structnames = {structname}
-        for structfield in idl.structs[structname].values():
-            structname2 = structfield.typename[3:]  # remove "GPU"
+        for attribute in idl.structs[structname].values():
+            structname2 = attribute.typename[3:]  # remove "GPU"
             if structname2 in idl.structs:
                 structnames.update(self._get_sub_structs(idl, structname2))
         return structnames

codegen/apiwriter.py

@@ -118,7 +118,7 @@ def write_structs():
         # Object-docstring as a comment
         for field in d.values():
             tp = idl.resolve_type(field.typename).strip("'")
-            if field.default is not None:
+            if field.default:
                 pylines.append(
                     resolve_crossrefs(f"#: * {field.name} :: {tp} = {field.default}")
                 )

codegen/idlparser.py

@@ -32,17 +32,32 @@ def get_idl_parser(*, allow_cache=True):
     return idl
 
 
-class StructField:
-    """A little object to specify the field of a struct."""
-
-    def __init__(self, line, name, typename, default=None):
-        self.line = line
-        self.name = name
-        self.typename = typename
+class Attribute:
+    """A little object to hold a function argument or struct field."""
+
+    def __init__(self, line):
+        self.line = line.strip().strip(",;").strip()
+
+        default = None  # None means 'no default' and "None" kinda means "auto".
+        arg = self.line
+        if "=" in arg:
+            arg, default = arg.rsplit("=", 1)
+            arg, default = arg.strip(), default.strip()
+        arg_type, arg_name = arg.strip().rsplit(" ", 1)
+        if arg_type.startswith("required "):
+            arg_type = arg_type.split(" ", 1)[1]
+            # required args should not have a default
+            assert default is None
+        elif arg_type.startswith("optional "):
+            arg_type = arg_type.split(" ", 1)[1]
+            default = default or "None"
+
+        self.name = arg_name
+        self.typename = arg_type
         self.default = default
 
     def __repr__(self):
-        return f"<StructField '{self.typename} {self.name}'>"
+        return f"<Attribute '{self.typename} {self.name}'>"
 
     def to_str(self):
         return self.line
@@ -65,7 +80,7 @@ class IdlParser:
 
     * flags: a dict mapping the (neutral) flag name to a dict of field-value pairs.
     * enums: a dict mapping the (Pythonic) enum name to a dict of field-value pairs.
-    * structs: a dict mapping the (Pythonic) struct name to a dict of StructField
+    * structs: a dict mapping the (Pythonic) struct name to a dict of Attribute
       objects.
     * classes: a dict mapping the (normalized) class name an Interface object.
 
@@ -199,6 +214,7 @@ def resolve_type(self, typename):
             "ImageBitmap": "memoryview",
             "ImageData": "memoryview",
             "VideoFrame": "memoryview",
+            "AllowSharedBufferSource": "memoryview",
             "GPUPipelineConstantValue": "float",
             "GPUExternalTexture": "object",
         }
@@ -208,22 +224,22 @@ def resolve_type(self, typename):
         if name.startswith("sequence<") and name.endswith(">"):
             name = name.split("<")[-1].rstrip(">")
             name = self.resolve_type(name).strip("'")
-            return f"'List[{name}]'"
+            return f"List[{name}]"
         elif name.startswith("record<") and name.endswith(">"):
             name = name.split("<")[-1].rstrip(">")
             names = [self.resolve_type(t).strip("'") for t in name.split(",")]
-            return f"'Dict[{', '.join(names)}]'"
+            return f"Dict[{', '.join(names)}]"
         elif " or " in name:
             name = name.strip("()")
             names = [self.resolve_type(t).strip("'") for t in name.split(" or ")]
             names = sorted(set(names))
-            return f"'Union[{', '.join(names)}]'"
+            return f"Union[{', '.join(names)}]"
 
         # Triage
         if name in __builtins__:
             return name  # ok
         elif name in self.classes:
-            return f"'{name}'"  # ok, but wrap in string because can be declared later
+            return name
         elif name.startswith("HTML"):
             return "object"  # anything, we ignore this stuff anyway
         elif name in ["OffscreenCanvas"]:
@@ -235,11 +251,11 @@ def resolve_type(self, typename):
             name = name[3:]
             name = name[:-4] if name.endswith("Dict") else name
             if name in self.flags:
-                return f"'flags.{name}'"
+                return f"flags.{name}"
             elif name in self.enums:
-                return f"'enums.{name}'"
+                return f"enums.{name}"
             elif name in self.structs:
-                return f"'structs.{name}'"
+                return f"structs.{name}"
             else:
                 # When this happens, update the code above or the pythonmap
                 raise RuntimeError("Encountered unknown IDL type: ", name)
@@ -373,19 +389,9 @@ def _parse(self):
                     if not line:
                         continue
                     assert line.endswith(";")
-                    arg = line.strip().strip(",;").strip()
-                    default = None
-                    if "=" in arg:
-                        arg, default = arg.rsplit("=", 1)
-                        arg, default = arg.strip(), default.strip()
-                    arg_type, arg_name = arg.strip().rsplit(" ", 1)
-                    if arg_type.startswith("required "):
-                        arg_type = arg_type[9:]
-                        # required args should not have a default
-                        assert default is None
-                    else:
-                        default = default or "None"
-                    d[arg_name] = StructField(line, arg_name, arg_type, default)
+
+                    attribute = Attribute(line)
+                    d[attribute.name] = attribute
                 self.structs[name] = d
             elif line.startswith(("[Exposed=", "[Serializable]")):
                 pass

pyproject.toml

@@ -76,7 +76,6 @@ ignore = [
     "B006",   # Do not use mutable data structures for argument defaults
     "B007",   # Loop control variable `x` not used within loop body
     "RUF012", # Mutable class attributes should be annotated with `typing.ClassVar`
-    "RUF013", # PEP 484 prohibits implicit `Optional` - we should fix this!
 ]
 
 [tool.coverage.report]