From 1499122ebf7281026f99946efdabbb916e2c6265 Mon Sep 17 00:00:00 2001
From: j3rome <jerome.abdelnour@usherbrooke.ca>
Date: Mon, 28 Jun 2021 14:52:09 -0400
Subject: [PATCH] * Changed how we scan for local module

+ Added option to override local module scan
---
 main.py  | 16 +++++++++++-----
 utils.py | 36 ++++++++++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/main.py b/main.py
index e8df2cb..7136cc4 100644
--- a/main.py
+++ b/main.py
@@ -7,7 +7,7 @@ from datetime import datetime
 from urllib.request import urlopen
 
 from utils import load_packages_from_requirements, get_mapping_files_from_pipreqs, user_response_multi_choices
-from utils import get_date_last_modified_python_file, get_python_filename_at_root, validate_cwd_is_git_repo
+from utils import get_date_last_modified_python_file, get_local_modules, validate_cwd_is_git_repo
 
 # TODO : Propose choice between date of first import or Added in requirements
 # TODO :    Other choices : When project was created, last commit (That wasnt on md file) get_date_last_modified_python_file()
@@ -24,12 +24,14 @@ from utils import get_date_last_modified_python_file, get_python_filename_at_roo
 
 # FIXME : Some unused imports might be important (Pillow for example)
 
+# LIMITATION : Might get conflicts with local imports
+
 EXTRACT_DATE_REGEX = re.compile(r'date\s-\s(\d+)')
 LETTER_REGEX = re.compile(r'[a-zA-Z]')
 
 parser = argparse.ArgumentParser("Python Requirements Version Guesser")
-parser.add_argument('--git_repo_path', type=str, default=None, required=False)  # TODO : CHDIR in this directory if provided
 parser.add_argument('--write', type=str, default=None, required=False, nargs='?', const='')
+parser.add_argument('--force_guess', type=str, default=None, required=False)
 
 
 def get_pypi_history(package_name, ignore_release_candidat=True):
@@ -230,12 +232,13 @@ if __name__ == "__main__":
     print("="*60)
     print("Python requirements guesser")
     print("="*60)
+    print(f"Guessing package versions for project '{os.getcwd()}'")
 
     if not validate_cwd_is_git_repo():
         print("[ERROR] py-reqs-guesser must be runned inside a git repository")
         exit(1)
 
-    print("\nFollow the steps to guess package versions based on when they were added to git")
+    print("Follow the steps to guess package versions based on when they were added to git.")
 
     args = parser.parse_args()
 
@@ -243,7 +246,10 @@ if __name__ == "__main__":
     stdlib_list, from_import_to_package_mapping, from_package_to_import_mapping = get_mapping_files_from_pipreqs()
 
     # Get local packages
-    local_packages = get_python_filename_at_root()
+    if args.force_guess:
+        args.force_guess = set(args.force_guess.strip().split(","))
+
+    local_packages = get_local_modules(print_modules=True, force_guess=args.force_guess)
 
     # Remove local_packages from the list of imports
     stdlib_list.update(local_packages)
@@ -269,7 +275,7 @@ if __name__ == "__main__":
     for package_name, version in sorted(packages, key=lambda x:x[0]):
         new_requirements_txt += f"{package_name}=={version}\n"
 
-    print("\n" + "="*60)
+    print("\n" + "="*60 + "\n")
     print("Requirements.txt :")
     print(new_requirements_txt)
     if args.write is None:
diff --git a/utils.py b/utils.py
index e5ea8d8..2ff2950 100644
--- a/utils.py
+++ b/utils.py
@@ -109,8 +109,40 @@ def load_packages_from_requirements(filepath):
     return packages
 
 
-def get_python_filename_at_root():
-    return [f[:-3] for f in os.listdir('.') if os.path.isfile(f) and f.endswith('.py')]
+def get_local_modules(print_modules=False, force_guess=None):
+    """
+    Gather list of the local python modules so we don't query pypi for those modules
+    Lets say we have the following file structure :
+        /project
+            - main.py
+            /utils
+                - common.py
+    common.py will be imported in main.py using 'from utils import common'
+    We therefore need to include the folder 'utils' in our exclusion list
+    """
+    if force_guess is None:
+        force_guess = set()
+
+    file_paths = subprocess.check_output('find . -name "*.py" -printf "%P\\n"', shell=True).decode().strip().split("\n")
+
+    modules = set()
+
+    for file_path in file_paths:
+        module = file_path.split('/')[0]
+        if '.py' in module:
+            module = module[:-3]
+
+        if module not in force_guess:
+            modules.add(module)
+
+    if print_modules:
+        print("\nWe detected the following local project modules :")
+        for module in modules:
+            print("    " + module)
+        print("We won't attempt to guess version for these packages (local files)")
+        print("In case of conflict, this can be overriden using --force_guess {package1},{package2},...")
+
+    return modules
 
 
 def get_date_last_modified_python_file():