Introduce gpustat package and restructure related codes

gpustat/__init__.py

@@ -0,0 +1,16 @@
+"""
+The gpustat module.
+"""
+
+__version__ = '0.5.0.dev1'
+
+from .gpustat import GPUStat, GPUStatCollection
+from .gpustat import new_query, print_gpustat
+from .gpustat import main
+
+
+__all__ = (
+    'GPUStat', 'GPUStatCollection',
+    'new_query', 'print_gpustat',
+    'main',
+)

gpustat/__main__.py

@@ -0,0 +1,7 @@
+try:
+    from .gpustat import main    # Try, python -m gpustat
+except ValueError:
+    raise RuntimeError("Attempted relative import in non-package.\n"
+                       " ... try 'python -m gpustat' instead?")
+
+main()

gpustat.py → gpustat/gpustat.py

@@ -1,13 +1,16 @@
 #!/usr/bin/env python
 
 """
-the gpustat script :)
+Implementation of gpustat
 
 @author Jongwook Choi
 @url https://github.com/wookayin/gpustat
 """
 
+from __future__ import absolute_import
+from __future__ import division
 from __future__ import print_function
+
 from datetime import datetime
 from six.moves import cStringIO as StringIO
 
@@ -22,7 +25,8 @@
 import pynvml as N
 from blessings import Terminal
 
-__version__ = '0.5.0.dev0'
+import sys
+from gpustat import __version__
 
 
 NOT_SUPPORTED = 'Not Supported'

test_gpustat.py → gpustat/test_gpustat.py

@@ -3,6 +3,7 @@
 """
 
 from __future__ import print_function
+from __future__ import absolute_import
 
 import gpustat
 
@@ -163,7 +164,18 @@ def remove_ansi_codes(s):
 class TestGPUStat(unittest.TestCase):
 
     @mock.patch('psutil.Process')
-    @mock.patch('gpustat.N')
+    @mock.patch('gpustat.gpustat.N')
+    def test_main(self, N, Process):
+        """
+        Test whether gpustat.main() works well. The behavior is mocked
+        exactly as in test_new_query_mocked().
+        """
+        _configure_mock(N, Process)
+        sys.argv = ['gpustat']
+        gpustat.main()
+
+    @mock.patch('psutil.Process')
+    @mock.patch('gpustat.gpustat.N')
     def test_new_query_mocked(self, N, Process):
         """
         A basic functionality test, in a case where everything is just normal.
@@ -185,7 +197,7 @@ def test_new_query_mocked(self, N, Process):
         self.assertEqual(unescaped, MOCK_EXPECTED_OUTPUT_FULL)
 
     @mock.patch('psutil.Process')
-    @mock.patch('gpustat.N')
+    @mock.patch('gpustat.gpustat.N')
     def test_new_query_mocked_nonexistent_pid(self, N, Process):
         """
         Test a case where nvidia query returns non-existent pids (see #16, #18)
@@ -196,7 +208,7 @@ def test_new_query_mocked_nonexistent_pid(self, N, Process):
         gpustats.print_formatted(fp=sys.stdout)
 
     @mock.patch('psutil.Process')
-    @mock.patch('gpustat.N')
+    @mock.patch('gpustat.gpustat.N')
     def test_attributes_and_items(self, N, Process):
         """
         Test whether each property of `GPUStat` instance is well-defined.
@@ -223,7 +235,7 @@ def test_attributes_and_items(self, N, Process):
 
     @unittest.skipIf(sys.version_info < (3, 4), "Only in Python 3.4+")
     @mock.patch('psutil.Process')
-    @mock.patch('gpustat.N')
+    @mock.patch('gpustat.gpustat.N')
     def test_args_endtoend(self, N, Process):
         """
         End-to-end testing given command line args.

setup.py

@@ -13,7 +13,7 @@ def read_readme():
 def read_version():
     # importing gpustat causes an ImportError :-)
     __PATH__ = os.path.abspath(os.path.dirname(__file__))
-    with open(os.path.join(__PATH__, 'gpustat.py')) as f:
+    with open(os.path.join(__PATH__, 'gpustat/__init__.py')) as f:
         version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
                                   f.read(), re.M)
     if version_match: