Ban onnx imports (#87)

In all places except for tests, serde, and selected passes and helpers.

Also cleaned up src/onnx_ir/_version_utils.py

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

pyproject.toml

@@ -107,11 +107,16 @@ ignore = [
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "pathlib".msg = "Using pathlib can impact performance. Use os.path instead"
+"onnx.helper".msg = "onnx helpers tend to be protobuf-y and slow. Consider using ir.tensor, ir.DataType and related methods instead"
+"onnx.numpy_helper".msg = "onnx numpy helpers tend to be slow. Consider using ir.tensor, ir.DataType and related methods instead"
+"onnx".msg = "Use onnx_ir methods and classes instead, or create an exception with `# noqa: TID251` if you need to use onnx directly"
 
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["TID252"] # Allow relative imports in init files
 "setup.py" = ["TID251"] # pathlib is allowed in supporting code
 "**/*_test.py" = ["TID251"] # pathlib is allowed in tests
+"tools/**/*.py" = ["TID251"] # OK to use in tools
+"src/onnx_ir/testing.py" = ["TID251"] # OK to use in tools
 
 [tool.ruff.lint.flake8-tidy-imports]
 # Disallow all relative imports.

src/onnx_ir/_convenience/__init__.py

@@ -19,7 +19,7 @@
 from collections.abc import Mapping, Sequence
 from typing import Union
 
-import onnx
+import onnx  # noqa: TID251
 
 from onnx_ir import _core, _enums, _protocols, serde, traversal
 

src/onnx_ir/_convenience/_constructors.py

@@ -13,7 +13,7 @@
 from collections.abc import Mapping, Sequence
 
 import numpy as np
-import onnx
+import onnx  # noqa: TID251
 
 from onnx_ir import _convenience, _core, _enums, _protocols, serde, tensor_adapters
 

src/onnx_ir/_io.py

@@ -9,7 +9,7 @@
 import os
 from typing import Callable
 
-import onnx
+import onnx  # noqa: TID251
 
 from onnx_ir import _core, _protocols, serde
 from onnx_ir import external_data as _external_data

src/onnx_ir/_version_utils.py

@@ -2,14 +2,15 @@
 # SPDX-License-Identifier: Apache-2.0
 """Version utils for testing."""
 
+# pylint: disable=import-outside-toplevel
 from __future__ import annotations
 
 import packaging.version
 
 
 def onnx_older_than(version: str) -> bool:
     """Returns True if the ONNX version is older than the given version."""
-    import onnx  # pylint: disable=import-outside-toplevel
+    import onnx  # noqa: TID251
 
     return (
         packaging.version.parse(onnx.__version__).release
@@ -19,50 +20,17 @@ def onnx_older_than(version: str) -> bool:
 
 def torch_older_than(version: str) -> bool:
     """Returns True if the torch version is older than the given version."""
-    import torch  # pylint: disable=import-outside-toplevel
+    import torch
 
     return (
         packaging.version.parse(torch.__version__).release
         < packaging.version.parse(version).release
     )
 
 
-def transformers_older_than(version: str) -> bool | None:
-    """Returns True if the transformers version is older than the given version."""
-    try:
-        import transformers  # pylint: disable=import-outside-toplevel
-    except ImportError:
-        return None
-
-    return (
-        packaging.version.parse(transformers.__version__).release
-        < packaging.version.parse(version).release
-    )
-
-
-def is_onnxruntime_training() -> bool:
-    """Returns True if the onnxruntime is onnxruntime-training."""
-    try:
-        from onnxruntime import training  # pylint: disable=import-outside-toplevel
-
-        assert training
-    except ImportError:
-        # onnxruntime not training
-        return False
-
-    try:
-        from onnxruntime.capi.onnxruntime_pybind11_state import (  # pylint: disable=import-outside-toplevel
-            OrtValueVector,
-        )
-    except ImportError:
-        return False
-
-    return hasattr(OrtValueVector, "push_back_batch")
-
-
 def onnxruntime_older_than(version: str) -> bool:
     """Returns True if the onnxruntime version is older than the given version."""
-    import onnxruntime  # pylint: disable=import-outside-toplevel
+    import onnxruntime
 
     return (
         packaging.version.parse(onnxruntime.__version__).release
@@ -72,20 +40,9 @@ def onnxruntime_older_than(version: str) -> bool:
 
 def numpy_older_than(version: str) -> bool:
     """Returns True if the numpy version is older than the given version."""
-    import numpy  # pylint: disable=import-outside-toplevel
+    import numpy
 
     return (
         packaging.version.parse(numpy.__version__).release
         < packaging.version.parse(version).release
     )
-
-
-def has_transformers():
-    """Tells if transformers is installed."""
-    try:
-        import transformers  # pylint: disable=import-outside-toplevel
-
-        assert transformers
-        return True  # noqa
-    except ImportError:
-        return False

src/onnx_ir/passes/common/_c_api_utils.py

@@ -10,7 +10,7 @@
 import onnx_ir as ir
 
 if TYPE_CHECKING:
-    import onnx
+    import onnx  # noqa: TID251
 
 
 logger = logging.getLogger(__name__)

src/onnx_ir/passes/common/onnx_checker.py

@@ -10,7 +10,7 @@
 
 from typing import Literal
 
-import onnx
+import onnx  # noqa: TID251
 
 import onnx_ir as ir
 from onnx_ir.passes.common import _c_api_utils

src/onnx_ir/passes/common/shape_inference.py

@@ -11,7 +11,7 @@
 
 import logging
 
-import onnx
+import onnx  # noqa: TID251
 
 import onnx_ir as ir
 from onnx_ir.passes.common import _c_api_utils

src/onnx_ir/passes/common/unused_removal.py

@@ -10,7 +10,7 @@
 
 import logging
 
-import onnx
+import onnx  # noqa: TID251
 
 import onnx_ir as ir
 

src/onnx_ir/serde.py

@@ -67,8 +67,8 @@
 from typing import Any, Callable
 
 import numpy as np
-import onnx
-import onnx.external_data_helper
+import onnx  # noqa: TID251
+import onnx.external_data_helper  # noqa: TID251
 
 from onnx_ir import _convenience, _core, _enums, _protocols, _type_casting