From 95d3a758a96faed071c6d452c4a13142ec837ff3 Mon Sep 17 00:00:00 2001 From: drbh Date: Thu, 7 Aug 2025 17:48:29 +0000 Subject: [PATCH 1/5] feat: add custom type for kernel modules --- src/kernels/utils.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/kernels/utils.py b/src/kernels/utils.py index f71855f..a7e996e 100644 --- a/src/kernels/utils.py +++ b/src/kernels/utils.py @@ -71,6 +71,17 @@ def universal_build_variant() -> str: return "torch-universal" +# Custom module type to identify dynamically loaded kernel modules. +# Using a subclass lets us distinguish these from regular imports. +class _KernelModule(ModuleType): + """Marker class for modules loaded dynamically from a path.""" + module_name: str + is_kernel: bool = True + + def __repr__(self): + return f"" + + def import_from_path(module_name: str, file_path: Path) -> ModuleType: # We cannot use the module name as-is, after adding it to `sys.modules`, # it would also be used for other imports. So, we make a module name that @@ -84,6 +95,8 @@ def import_from_path(module_name: str, file_path: Path) -> ModuleType: module = importlib.util.module_from_spec(spec) if module is None: raise ImportError(f"Cannot load module {module_name} from spec") + module.__class__ = _KernelModule + module.module_name = module_name sys.modules[module_name] = module spec.loader.exec_module(module) # type: ignore return module From 53723ea9869c90e1ddc43b46fba091f4b16101ff Mon Sep 17 00:00:00 2001 From: drbh Date: Thu, 7 Aug 2025 17:53:51 +0000 Subject: [PATCH 2/5] fix: improve name --- src/kernels/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/kernels/utils.py b/src/kernels/utils.py index a7e996e..598bb26 100644 --- a/src/kernels/utils.py +++ b/src/kernels/utils.py @@ -73,7 +73,7 @@ def universal_build_variant() -> str: # Custom module type to identify dynamically loaded kernel modules. # Using a subclass lets us distinguish these from regular imports. -class _KernelModule(ModuleType): +class _KernelModuleType(ModuleType): """Marker class for modules loaded dynamically from a path.""" module_name: str is_kernel: bool = True @@ -95,7 +95,7 @@ def import_from_path(module_name: str, file_path: Path) -> ModuleType: module = importlib.util.module_from_spec(spec) if module is None: raise ImportError(f"Cannot load module {module_name} from spec") - module.__class__ = _KernelModule + module.__class__ = _KernelModuleType module.module_name = module_name sys.modules[module_name] = module spec.loader.exec_module(module) # type: ignore From 338752c367d508f53e373ec0f20c6aeac7139acb Mon Sep 17 00:00:00 2001 From: drbh Date: Thu, 7 Aug 2025 17:59:32 +0000 Subject: [PATCH 3/5] fix: improve metaclass repr --- src/kernels/utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/kernels/utils.py b/src/kernels/utils.py index 598bb26..a1027cb 100644 --- a/src/kernels/utils.py +++ b/src/kernels/utils.py @@ -71,10 +71,17 @@ def universal_build_variant() -> str: return "torch-universal" +# Metaclass to allow overriding the `__repr__` method for kernel modules. +class _KernelModuleMeta(type): + def __repr__(self): + return "" + + # Custom module type to identify dynamically loaded kernel modules. # Using a subclass lets us distinguish these from regular imports. -class _KernelModuleType(ModuleType): +class _KernelModuleType(ModuleType, metaclass=_KernelModuleMeta): """Marker class for modules loaded dynamically from a path.""" + module_name: str is_kernel: bool = True From e487945cf6a39525d48c907b86dbaa2bd770d869 Mon Sep 17 00:00:00 2001 From: drbh Date: Thu, 7 Aug 2025 18:14:07 +0000 Subject: [PATCH 4/5] fix: help mypy type checking --- src/kernels/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/kernels/utils.py b/src/kernels/utils.py index a1027cb..0fe9b63 100644 --- a/src/kernels/utils.py +++ b/src/kernels/utils.py @@ -103,6 +103,7 @@ def import_from_path(module_name: str, file_path: Path) -> ModuleType: if module is None: raise ImportError(f"Cannot load module {module_name} from spec") module.__class__ = _KernelModuleType + assert isinstance(module, _KernelModuleType) # for mypy type checking module.module_name = module_name sys.modules[module_name] = module spec.loader.exec_module(module) # type: ignore From e7173d9a280581b66c539ca170ada0e99edb4080 Mon Sep 17 00:00:00 2001 From: drbh Date: Thu, 7 Aug 2025 18:16:34 +0000 Subject: [PATCH 5/5] fix: lint correction --- src/kernels/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/kernels/utils.py b/src/kernels/utils.py index 0fe9b63..3a0e89b 100644 --- a/src/kernels/utils.py +++ b/src/kernels/utils.py @@ -103,7 +103,7 @@ def import_from_path(module_name: str, file_path: Path) -> ModuleType: if module is None: raise ImportError(f"Cannot load module {module_name} from spec") module.__class__ = _KernelModuleType - assert isinstance(module, _KernelModuleType) # for mypy type checking + assert isinstance(module, _KernelModuleType) # for mypy type checking module.module_name = module_name sys.modules[module_name] = module spec.loader.exec_module(module) # type: ignore