Merge pull request #3 from littmus/master

Fixed problems with local imports and imports with 'as'

pipreqs/pipreqs.py

@@ -33,12 +33,12 @@ def get_all_imports(start_path):
     logging.debug('Traversing tree, start: %s', start_path)
     for root, dirs, files in os.walk(start_path):
         packages.append(os.path.basename(root))
+        files = [fn for fn in files if os.path.splitext(fn)[1] == ".py"]
+        packages += [os.path.splitext(fn)[0] for fn in files]
         for file_name in files:
-            if file_name[-3:] != ".py":
-                continue
-
             with open(os.path.join(root, file_name), "r") as file_object:
-                for line in file_object:
+                lines = filter(lambda l:len(l) > 0, map(lambda l:l.strip(), file_object))
+                for line in lines:
                     if line[0] == "#":
                         continue
                     if "(" in line:
@@ -51,6 +51,9 @@ def get_all_imports(start_path):
                             if "," in item:
                                 for match in item.split(","):
                                     imports.append(match.strip())
+                            elif " as " in item:
+                                to_append = item.split(" as ")[0]
+                                imports.append(to_append.strip())
                             else:
                                 to_append = item if "." not in item else item.split(".")[0]
                                 imports.append(to_append.strip())
@@ -109,4 +112,4 @@ def main():  # pragma: no cover
 
 
 if __name__ == '__main__':
-    main()  # pragma: no cover
+    main()  # pragma: no cover

tests/_data/test.py

@@ -10,11 +10,15 @@ import sys
 import signal
 import requests
 import nonexistendmodule
-# Ignore this Line
+# import django
 import flask.ext.somext
 from sqlalchemy import model
+try:
+    import ujson as json
+except ImportError:
+    import json
 
-
+import models
 
 def main():
 	pass

tests/test_pipreqs.py

@@ -15,25 +15,27 @@ from pipreqs import pipreqs
 
 class TestPipreqs(unittest.TestCase):
     def setUp(self):
-        self.modules = ['flask', 'requests', 'sqlalchemy', 'docopt', 'nonexistendmodule']
+        self.modules = ['flask', 'requests', 'sqlalchemy', 'docopt', 'ujson', 'nonexistendmodule']
         self.project = os.path.join(os.path.dirname(__file__), "_data")
         self.requirements_path = os.path.join(self.project, "requirements.txt")
 
     def test_get_all_imports(self):
         imports = pipreqs.get_all_imports(self.project)
-        self.assertEqual(len(imports), 5, "Incorrect Imports array length")
+        self.assertEqual(len(imports), 6, "Incorrect Imports array length")
         for item in imports:
             self.assertTrue(item in self.modules, "Import is missing")
         self.assertFalse("time" in imports)
         self.assertFalse("logging" in imports)
         self.assertFalse("curses" in imports)
         self.assertFalse("__future__" in imports)
+        self.assertFalse("django" in imports)
+        self.assertFalse("models" in imports)
 
     def test_get_imports_info(self):
         imports = pipreqs.get_all_imports(self.project)
         with_info = pipreqs.get_imports_info(imports)
-        # Should contain only 4 Elements without the "nonexistendmodule"
-        self.assertEqual(len(with_info), 4, "Length of imports array with info is wrong")
+        # Should contain only 5 Elements without the "nonexistendmodule"
+        self.assertEqual(len(with_info), 5, "Length of imports array with info is wrong")
         for item in with_info:
             self.assertTrue(item['name'] in self.modules, "Import item appears to be missing")