[fix] Python2 import fix

dvolgyes · Jun 25, 2018 · 3eb0181 · 3eb0181
1 parent 676fff5
commit 3eb0181
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 import setuptools
-import src.zenodo_get as zget
+from src import zenodo_get as zget
 
 setuptools.setup(
     name=zget.__title__,

diff --git a/src/__init__.py b/src/__init__.py
diff --git a/src/zenodo_get.py b/src/zenodo_get.py
@@ -94,7 +94,9 @@ def check_hash(filename, checksum):
     return value, digest
 
 
-if __name__ == '__main__':
+def zenodo_get(argv=None):
+    if argv is None:
+        argv = sys.argv[1:]
 
     parser = OptionParser(
         usage='%prog [options] RECORD_OR_DOI', version=f'%prog {__version__}'
@@ -207,7 +209,7 @@ def check_hash(filename, checksum):
         default=15.,
     )
 
-    (options, args) = parser.parse_args()
+    (options, args) = parser.parse_args(argv)
 
     if options.cite:
         print('Reference for this software:')
@@ -219,8 +221,7 @@ def check_hash(filename, checksum):
 
     if len(args) > 0:
         try:
-            t = int(args[0])
-            options.record = args[0]
+            options.record = str(int(args[0]))
         except ValueError:
             options.doi = args[0]
     elif options.doi is None and options.record is None:
@@ -341,3 +342,7 @@ def check_hash(filename, checksum):
     else:
         eprint('Record could not get accessed.')
         sys.exit(1)
+
+
+if __name__ == '__main__':
+    zenodo_get()