diff --git a/qcog_python_client/qcog/pytorch/validate/utils.py b/qcog_python_client/qcog/pytorch/validate/utils.py index 0bd1756..f0f0287 100644 --- a/qcog_python_client/qcog/pytorch/validate/utils.py +++ b/qcog_python_client/qcog/pytorch/validate/utils.py @@ -89,9 +89,12 @@ def get_third_party_imports(source_code: io.BytesIO, package_path: str) -> set[s print("------> Python Sys Lib: ", python_sys_lib) print("------> Common path: ", common_path) - if common_path == python_sys_lib: + if infolder(spec.origin, python_sys_lib): + # Check if the module is directly inside the folder + # or if it's a subpackage continue + print(f"------> Third Party Package: {base_package}") third_party_packages.add(base_package) return third_party_packages @@ -102,3 +105,31 @@ def is_package_module(module_path: str) -> bool: module_path = module_path if module_path.endswith(".py") else module_path + ".py" return os.path.isfile(module_path) + + +def infolder(module_path: str, folder_path: str) -> bool: + """Check if a module is in a folder.""" + # Folder Path should be a substring of the module path + if not module_path.startswith(folder_path): + return False + + # Check if the file that the module is pointing at is a `__init__.py` file. + # If it is then it's a package and the previous segment of the path + # is the package name, otherwise the module is the package name. + + package_name = None + filename = os.path.basename(module_path) + if filename == "__init__.py": + package_name = os.path.basename(os.path.dirname(module_path)) + else: + package_name = os.path.basename(module_path) + + print("Package Name: ", package_name) + # The module is in the folder if the `package_name` is right after + # the `folder_path` in the `module_path` + # so something like `/` should be + # the a substring of the `module_path` + module_subpath_candidate = os.path.join(folder_path, package_name) + + return module_path.startswith(module_subpath_candidate) +