[UR] Cleanup Python linter issues add type hints (#18694)

Continue cleanup of the Python generator scripts, resolving Pyright
issues and adding typing hints. This should enable easier refactoring of
the scripts in future commits.

Author: kbenzie

unified-runtime/scripts/generate_code.py

@@ -5,7 +5,6 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import os
-import re
 import util
 
 
@@ -428,7 +427,7 @@ def _mako_interface_loader_api(
     template = f"ur_interface_loader.{ext}.mako"
     fin = os.path.join(templates_dir, template)
 
-    name = f"ur_interface_loader"
+    name = "ur_interface_loader"
 
     filename = f"{name}.{ext}"
     fout = os.path.join(dstpath, filename)

unified-runtime/scripts/generate_ids.py

@@ -6,7 +6,9 @@
 
 """Generates a unique id for each spec function that doesn't have it."""
 
-from fileinput import FileInput
+from typing import Callable, List
+from yaml.dumper import Dumper
+from yaml.representer import Node
 import util
 import yaml
 import re
@@ -15,46 +17,48 @@
 ENUM_NAME = "$x_function_t"
 
 
-class quoted(str):
+class Quoted(str):
     pass
 
 
-def quoted_presenter(dumper, data):
+def quoted_presenter(dumper: Dumper, data: Quoted) -> Node:
     return dumper.represent_scalar("tag:yaml.org,2002:str", data, style='"')
 
 
-def get_registry_header():
+def get_registry_header() -> dict:
     return {
         "type": "header",
-        "desc": quoted("Intel $OneApi Unified Runtime function registry"),
-        "ordinal": quoted(-1),
+        "desc": Quoted("Intel $OneApi Unified Runtime function registry"),
+        "ordinal": Quoted(-1),
     }
 
 
-def write_registry(data, path):
+def write_registry(data: List[dict], path: str) -> None:
     with open(path, "w") as fout:
-        yaml.add_representer(quoted, quoted_presenter)
+        yaml.add_representer(Quoted, quoted_presenter)
         yaml.dump_all(
             data, fout, default_flow_style=False, sort_keys=False, explicit_start=True
         )
 
 
-def find_type_in_specs(specs, type):
+def find_type_in_specs(specs: List[dict], type: str) -> dict:
     return [obj for s in specs for obj in s["objects"] if obj["name"] == type][0]
 
 
-def get_max_enum(enum):
+def get_max_enum(enum: dict) -> int:
     return int(max(enum["etors"], key=lambda x: int(x["value"]))["value"])
 
 
-def copy_and_strip_prefix_from_enums(enum, prefix):
+def copy_and_strip_prefix_from_enums(enum: dict, prefix: str) -> dict:
     cpy = copy.deepcopy(enum)
     for etor in cpy["etors"]:
         etor["name"] = etor["name"][len(prefix) :]
     return cpy
 
 
-def generate_function_type(specs, meta, update_fn) -> dict:
+def generate_function_type(
+    specs: List[dict], meta: dict, update_fn: Callable[[dict, dict], None]
+) -> dict:
     existing_function_type = find_type_in_specs(specs, "$x_function_t")
     existing_etors = {
         etor["name"]: etor["value"] for etor in existing_function_type["etors"]
@@ -84,7 +88,9 @@ def generate_function_type(specs, meta, update_fn) -> dict:
     return copy_and_strip_prefix_from_enums(existing_function_type, "$X_FUNCTION_")
 
 
-def generate_structure_type(specs, meta, refresh_fn) -> dict:
+def generate_structure_type(
+    specs: List[dict], meta: dict, refresh_fn: Callable[[dict, dict], None]
+) -> dict:
     structure_type = find_type_in_specs(specs, "$x_structure_type_t")
     extended_structs = [
         obj
@@ -124,7 +130,9 @@ def generate_structure_type(specs, meta, refresh_fn) -> dict:
     return copy_and_strip_prefix_from_enums(structure_type, "$X_STRUCTURE_TYPE_")
 
 
-def generate_registry(path, specs, meta, update_fn):
+def generate_registry(
+    path: str, specs: List[dict], meta: dict, update_fn: Callable[[dict, dict], None]
+) -> None:
     try:
         write_registry(
             [

unified-runtime/scripts/parse_specs.py

@@ -14,7 +14,7 @@
 from templates.helper import param_traits, type_traits, value_traits
 import ctypes
 import itertools
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 from version import Version
 
 
@@ -86,7 +86,9 @@ def _validate_doc(f, d, tags, line_num, meta):
     validate documents meet some basic (easily detectable) requirements of code
     generation
     """
-    is_iso = lambda x: re.match(r"[_a-zA-Z][_a-zA-Z0-9]{0,30}", x)
+
+    def is_iso(x):
+        return re.match(r"[_a-zA-Z][_a-zA-Z0-9]{0,30}", x)
 
     def __validate_ordinal(d):
         if "ordinal" in d:
@@ -95,7 +97,7 @@ def __validate_ordinal(d):
 
             try:
                 ordinal = str(int(d["ordinal"]))
-            except:
+            except BaseException:
                 ordinal = None
 
             if ordinal != d["ordinal"]:
@@ -110,7 +112,7 @@ def __validate_version(d, prefix="", base_version=default_version):
 
             try:
                 version = str(d["version"])
-            except:
+            except BaseException:
                 version = None
 
             if version != d["version"]:
@@ -124,7 +126,7 @@ def __validate_tag(d, key, tags, case):
                 return x
         return None
 
-    def __validate_desc(desc):
+    def __validate_desc(desc: Union[dict, str], prefix: str):
         if isinstance(desc, dict):
             for k, v in desc.items():
                 if not isinstance(k, str):
@@ -134,7 +136,7 @@ def __validate_desc(desc):
 
                 try:
                     version = str(k)
-                except:
+                except BaseException:
                     version = None
 
                 if version != k:
@@ -244,34 +246,36 @@ def __validate_etors(d, tags):
         for i, item in enumerate(d["etors"]):
             prefix = "'etors'[%s] " % i
             if not isinstance(item, dict):
-                raise Exception(prefix + "must be a mapping: '%s'" % (i, type(item)))
+                raise Exception(
+                    prefix + "must be a mapping: %d, '%s'" % (i, type(item))
+                )
 
             if ("desc" not in item) or ("name" not in item):
                 raise Exception(
                     prefix + "requires the following scalar fields: {`desc`, `name`}"
                 )
 
-            if "extend" in d and d.get("extend") == True and "value" not in item:
+            if "extend" in d and d.get("extend") and "value" not in item:
                 raise Exception(
                     prefix
                     + "must include a value for experimental features: {`value`: `0xabcd`}"
                 )
 
             if typed:
-                type = extract_type(item["desc"])
-                if type is None:
+                ty = extract_type(item["desc"])
+                if ty is None:
                     raise Exception(
                         prefix
                         + "typed etor "
                         + item["name"]
                         + " must begin with a type identifier: [type]"
                     )
-                type_name = _subt(type, tags)
+                type_name = _subt(ty, tags)
                 if not is_iso(type_name):
                     raise Exception(
                         prefix
                         + "type "
-                        + str(type)
+                        + str(ty)
                         + " in a typed etor "
                         + item["name"]
                         + " must be a valid ISO C identifier"
@@ -331,7 +335,9 @@ def has_handle(members, meta):
                 if type_traits.is_handle(m):
                     return True
                 if type_traits.is_struct(m, meta):
-                    return has_handle(type_traits.get_struct_members(m["type"]), meta)
+                    return has_handle(
+                        type_traits.get_struct_members(m["type"], meta), meta
+                    )
             return False
 
         for m in members:
@@ -350,7 +356,7 @@ def has_handle(members, meta):
                 # exception messages.
                 __validate_struct_range_members(name, member_members, meta)
 
-    def __validate_members(d, tags, meta):
+    def __validate_members(d, meta):
         if "members" not in d:
             raise Exception(
                 "'%s' requires the following sequence of mappings: {`members`}"
@@ -365,15 +371,17 @@ def __validate_members(d, tags, meta):
         for i, item in enumerate(d["members"]):
             prefix = "'members'[%s] " % i
             if not isinstance(item, dict):
-                raise Exception(prefix + "must be a mapping: '%s'" % (i, type(item)))
+                raise Exception(
+                    prefix + "must be a mapping: %d, '%s'" % (i, type(item))
+                )
 
             if ("desc" not in item) or ("type" not in item) or ("name" not in item):
                 raise Exception(
                     prefix
                     + "requires the following scalar fields: {`desc`, 'type', `name`}"
                 )
 
-            annotation = __validate_desc(item["desc"])
+            annotation = __validate_desc(item["desc"], prefix)
             if not annotation:
                 raise Exception(
                     prefix + "'desc' must start with {'[in]', '[out]', '[in,out]'}"
@@ -409,7 +417,7 @@ def __validate_members(d, tags, meta):
                 )
             max_ver = ver
 
-    def __validate_params(d, tags, meta):
+    def __validate_params(d, meta):
         if "params" not in d:
             raise Exception(
                 "'function' requires the following sequence of mappings: {`params`}"
@@ -420,19 +428,25 @@ def __validate_params(d, tags, meta):
 
         d_ver = Version(d.get("version", default_version))
         max_ver = d_ver
-        min = {"[in]": None, "[out]": None, "[in,out]": None}
+        min: Dict[str, Union[int, None]] = {
+            "[in]": None,
+            "[out]": None,
+            "[in,out]": None,
+        }
         for i, item in enumerate(d["params"]):
             prefix = "'params'[%s] " % i
             if not isinstance(item, dict):
-                raise Exception(prefix + "must be a mapping: '%s'" % (i, type(item)))
+                raise Exception(
+                    prefix + "must be a mapping: %d, '%s'" % (i, type(item))
+                )
 
             if ("desc" not in item) or ("type" not in item) or ("name" not in item):
                 raise Exception(
                     prefix
                     + "requires the following scalar fields: {`desc`, 'type', `name`}"
                 )
 
-            annotation = __validate_desc(item["desc"])
+            annotation = __validate_desc(item["desc"], prefix)
             if not annotation:
                 raise Exception(
                     prefix + "'desc' must start with {'[in]', '[out]', '[in,out]'}"
@@ -576,7 +590,7 @@ def __validate_union_tag(d):
                 __validate_union_tag(d)
             __validate_type(d, "name", tags)
             __validate_base(d)
-            __validate_members(d, tags, meta)
+            __validate_members(d, meta)
             __validate_details(d)
             __validate_ordinal(d)
             __validate_version(d)
@@ -594,7 +608,7 @@ def __validate_union_tag(d):
             else:
                 __validate_name(d, "name", tags, case="camel")
 
-            __validate_params(d, tags, meta)
+            __validate_params(d, meta)
             __validate_details(d)
             __validate_ordinal(d)
             __validate_version(d)
@@ -734,6 +748,7 @@ def _generate_meta(d, ordinal, meta):
         if "enum" == type:
             value = -1
             max_value = -1
+            max_index = -1
             bit_mask = 0
             meta[type][name]["etors"] = []
             for idx, etor in enumerate(d["etors"]):
@@ -902,7 +917,8 @@ def append_nullchecks(param, accessor: str):
                 )
 
         def append_enum_checks(param, accessor: str):
-            ptypename = type_traits.base(param["type"])
+            typename = type_traits.base(param["type"])
+            assert typename
 
             prefix = "`"
             if param_traits.is_optional(item):
@@ -924,20 +940,19 @@ def append_enum_checks(param, accessor: str):
             else:
                 if (
                     type_traits.is_flags(param["type"])
-                    and "bit_mask" in meta["enum"][ptypename].keys()
+                    and "bit_mask" in meta["enum"][typename].keys()
                 ):
                     _append(
                         rets,
                         "$X_RESULT_ERROR_INVALID_ENUMERATION",
                         prefix
-                        + "%s & %s`" % (ptypename.upper()[:-2] + "_MASK", accessor),
+                        + "%s & %s`" % (typename.upper()[:-2] + "_MASK", accessor),
                     )
                 else:
                     _append(
                         rets,
                         "$X_RESULT_ERROR_INVALID_ENUMERATION",
-                        prefix
-                        + "%s < %s`" % (meta["enum"][ptypename]["max"], accessor),
+                        prefix + "%s < %s`" % (meta["enum"][typename]["max"], accessor),
                     )
 
         # generate results based on parameters
@@ -958,7 +973,7 @@ def append_enum_checks(param, accessor: str):
             ):
                 typename = type_traits.base(item["type"])
                 # walk each entry in the desc for pointers and enums
-                for i, m in enumerate(meta["struct"][typename]["members"]):
+                for m in meta["struct"][typename]["members"]:
                     if param_traits.is_nocheck(m):
                         continue
 
@@ -1050,8 +1065,9 @@ def _validate_ext_enum_range(extension, enum) -> bool:
             if value in existing_values:
                 return False
             return True
-    except:
-        return False
+    except BaseException:
+        pass
+    return False
 
 
 def _extend_enums(enum_extensions, specs, meta):
@@ -1068,9 +1084,9 @@ def _extend_enums(enum_extensions, specs, meta):
             for obj in s["objects"]
             if obj["type"] == "enum" and k == obj["name"] and not obj.get("extend")
         ][0]
-        for i, extension in enumerate(group):
+        for extension in group:
             if not _validate_ext_enum_range(extension, matching_enum):
-                raise Exception(f"Invalid enum values.")
+                raise Exception("Invalid enum values.")
             matching_enum["etors"].extend(extension["etors"])
 
         _refresh_enum_meta(matching_enum, meta)
@@ -1118,7 +1134,7 @@ def parse(section, version, tags, meta, ref):
             if not d:
                 continue
 
-            if d["type"] == "enum" and d.get("extend") == True:
+            if d["type"] == "enum" and d.get("extend"):
                 # enum extensions are resolved later
                 enum_extensions.append(d)
                 continue