fix build process

Author: SkafteNicki

setup.py

@@ -39,7 +39,7 @@ def run(self):
 
 
 setup(
-    name="torchplot",
+    name=torchplot.__name__,
     version=torchplot.__version__,
     description=torchplot.__docs__,
     long_description=long_description,
@@ -65,6 +65,7 @@ def run(self):
         "Operating System :: OS Independent",
         # Specify the Python versions you support here. In particular, ensure
         # that you indicate whether you support Python 2, Python 3 or both.
-        "Programming Language :: Python :: 3"
+        "Programming Language :: Python :: 3",
     ],
+    cmdclass={"clean": CleanCommand},
 )

tests/test_torchplot.py

@@ -58,7 +58,7 @@ def test_cpu(test_case):
 def test_gpu(test_case):
     """ test that it works on gpu """
     assert tp.plot(
-        test_case.x.cuda() if isinstance(test_case.x, torch.Tensor) else test_case.x, 
-        test_case.y.cuda() if isinstance(test_case.y, torch.Tensor) else test_case.y, 
-        "."
+        test_case.x.cuda() if isinstance(test_case.x, torch.Tensor) else test_case.x,
+        test_case.y.cuda() if isinstance(test_case.y, torch.Tensor) else test_case.y,
+        ".",
     )

torchplot/__init__.py

@@ -1,19 +1,31 @@
 #!/usr/bin/env python
-"""Root package info."""
 import os
 import time
 
+__name__ = "torchplot"
 _this_year = time.strftime("%Y")
 __version__ = "0.1.5"
 __author__ = "Nicki Skafte Detlefsen et al."
 __author_email__ = "[email protected]"
 __license__ = "Apache-2.0"
 __copyright__ = f"Copyright (c) 2018-{_this_year}, {__author__}."
 __homepage__ = "https://github.com/CenterBioML/torchplot"
-
 __docs__ = "Plotting pytorch tensors made easy"
 
 PACKAGE_ROOT = os.path.dirname(__file__)
 PROJECT_ROOT = os.path.dirname(PACKAGE_ROOT)
 
-from .core import *
+try:
+    # This variable is injected in the __builtins__ by the build process
+    _ = None if __TORCHPLOT_SETUP__ else None
+except NameError:
+    __TORCHPLOT_SETUP__: bool = False
+
+
+if __TORCHPLOT_SETUP__:
+    import sys  # pragma: no-cover
+
+    sys.stdout.write(f"Partial import of `{__name__}` during the build process.\n")  # pragma: no-cover
+    # We are not importing the rest of the package during the build process, as it may not be compiled yet
+else:
+    from .core import *