[V2 Logger] utils for import

src/deepsparse/loggers_v2/utils.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+import re
+from typing import Any, Type
+
+
+LOGGER_REGISTRY = "src.deepsparse.loggers_v2.registry.__init__"
+
+
+def import_from_registry(name: str) -> Type[Any]:
+    """
+    Import `name` from the LOGGER_REGISTRY
+
+    :param name: name of the function or class name in LOGGER_REGISTRY
+    :return: Function or class object
+    """
+    module = importlib.import_module(LOGGER_REGISTRY)
+    try:
+        return getattr(module, name)
+    except AttributeError:
+        raise AttributeError(
+            f"Cannot import class/func with name '{name}' from {LOGGER_REGISTRY}"
+        )
+
+
+def import_from_path(path: str) -> Type[Any]:
+    """
+    Import the module and the name of the function/class separated by :
+
+    Examples:
+      path = "/path/to/file.py:func_name"
+      path = "/path/to/file:class_name"
+
+    :param path: path including the file path and object name
+    :return Function or class object
+
+    """
+    path, class_name = path.split(":")
+    _path = path
+
+    path = path.split(".py")[0]
+    path = re.sub(r"/+", ".", path)
+    try:
+        module = importlib.import_module(path)
+    except ImportError:
+        raise ImportError(f"Cannot find module with path {_path}")
+
+    try:
+        return getattr(module, class_name)
+    except AttributeError:
+        raise AttributeError(f"Cannot find {class_name} in {_path}")

tests/deepsparse/loggers_v2/test_util.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+from deepsparse.loggers_v2.utils import (
+    LOGGER_REGISTRY,
+    import_from_path,
+    import_from_registry,
+)
+
+
+@pytest.mark.parametrize(
+    "name, is_successful",
+    [
+        ("PythonLogger", True),
+        ("max", True),
+        ("blah", False),
+    ],
+)
+def test_import_from_registry(name, is_successful):
+    if is_successful:
+        assert import_from_registry(name) is not None
+    else:
+        with pytest.raises(AttributeError):
+            import_from_registry(name)
+
+
+@pytest.mark.parametrize(
+    "path, is_successful",
+    [
+        (f"{LOGGER_REGISTRY}.py:PythonLogger", True),
+        (f"{LOGGER_REGISTRY}:PythonLogger", True),
+        ("foo/bar:blah", False),
+        (f"{LOGGER_REGISTRY}:blah", False),
+    ],
+)
+def test_import_from_path(path, is_successful):
+    if is_successful:
+        assert import_from_path(path) is not None
+    else:
+        with pytest.raises((AttributeError, ImportError)):
+            import_from_path(path)