get installed modules

Author: KRRT7

codeflash/code_utils/code_utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import importlib.metadata
 import os
 import re
 import shutil
@@ -10,6 +11,7 @@
 from functools import lru_cache
 from pathlib import Path
 from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING
 
 import tomlkit
 
@@ -19,8 +21,12 @@
 ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
 
 
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+
 @contextmanager
-def custom_addopts() -> None:
+def custom_addopts() -> Generator[None, None, None]:
     pyproject_file = find_pyproject_toml()
     original_content = None
     non_blacklist_plugin_args = ""
@@ -58,7 +64,7 @@ def custom_addopts() -> None:
 
 
 @contextmanager
-def add_addopts_to_pyproject() -> None:
+def add_addopts_to_pyproject() -> Generator[None, None, None]:
     pyproject_file = find_pyproject_toml()
     original_content = None
     try:
@@ -220,3 +226,32 @@ def exit_with_message(message: str, *, error_on_exit: bool = False) -> None:
     paneled_text(message, panel_args={"style": "red"})
 
     sys.exit(1 if error_on_exit else 0)
+
+
+blacklist_installed_pkgs = {
+    "codeflash",
+    "pytest",
+    "coverage",
+    "__",  # this is for private packages  or ones that contain "__" in order to mangle names i.e 3204bda914b7f2c6f497__mypyc
+    "setuptools",
+    "pip",
+    "wheel",
+    "importlib_metadata",
+    "importlib_resources",
+    "isort",
+    "black",
+    "tomlkit",
+    "stubs",
+}
+
+
+def get_installed_packages() -> set[str]:
+    pkgs = importlib.metadata.packages_distributions().keys()
+    return {
+        pkg
+        for pkg in pkgs
+        if not any(blacklisted in pkg for blacklisted in blacklist_installed_pkgs) and not pkg.startswith("_")
+    }
+
+
+print(get_installed_packages())