Redirect types and datamodel modules

Author: gmarkall

numba_cuda/numba/cuda/datamodel/cuda_manager.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import weakref
+from collections import ChainMap
+
+from numba.cuda import types
+
+
+class DataModelManager(object):
+    """Manages mapping of FE types to their corresponding data model"""
+
+    def __init__(self, handlers=None):
+        """
+        Parameters
+        -----------
+        handlers: Mapping[Type, DataModel] or None
+            Optionally provide the initial handlers mapping.
+        """
+        # { numba type class -> model factory }
+        self._handlers = handlers or {}
+        # { numba type instance -> model instance }
+        self._cache = weakref.WeakKeyDictionary()
+
+    def register(self, fetypecls, handler):
+        """Register the datamodel factory corresponding to a frontend-type class"""
+        assert issubclass(fetypecls, types.Type)
+        self._handlers[fetypecls] = handler
+
+    def lookup(self, fetype):
+        """Returns the corresponding datamodel given the frontend-type instance"""
+        try:
+            return self._cache[fetype]
+        except KeyError:
+            pass
+        handler = self._handlers[type(fetype)]
+        model = self._cache[fetype] = handler(self, fetype)
+        return model
+
+    def __getitem__(self, fetype):
+        """Shorthand for lookup()"""
+        return self.lookup(fetype)
+
+    def copy(self):
+        """
+        Make a copy of the manager.
+        Use this to inherit from the default data model and specialize it
+        for custom target.
+        """
+        return DataModelManager(self._handlers.copy())
+
+    def chain(self, other_manager):
+        """Create a new DataModelManager by chaining the handlers mapping of
+        `other_manager` with a fresh handlers mapping.
+
+        Any existing and new handlers inserted to `other_manager` will be
+        visible to the new manager. Any handlers inserted to the new manager
+        can override existing handlers in `other_manager` without actually
+        mutating `other_manager`.
+
+        Parameters
+        ----------
+        other_manager: DataModelManager
+        """
+        chained = ChainMap(self._handlers, other_manager._handlers)
+        return DataModelManager(chained)