diff --git a/pipreqs/pipreqs.py b/pipreqs/pipreqs.py index 7a63556..c8f8d6d 100755 --- a/pipreqs/pipreqs.py +++ b/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 \ No newline at end of file + main() # pragma: no cover diff --git a/tests/_data/models.py b/tests/_data/models.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/_data/test.py b/tests/_data/test.py index 03b688c..b033ef7 100644 --- a/tests/_data/test.py +++ b/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 diff --git a/tests/test_pipreqs.py b/tests/test_pipreqs.py index f6dc73e..a2001dd 100755 --- a/tests/test_pipreqs.py +++ b/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")