204 lines
4.1 KiB
Python
204 lines
4.1 KiB
Python
import ast
|
|
import os
|
|
import sys
|
|
|
|
from .abstract_check import ASTChecker # pyre-ignore[21]
|
|
|
|
# Allowed packages list from original check_dependencies.py
|
|
ALLOWED_PACKAGES = {
|
|
# Python Standard Library (Major ones)
|
|
"ast",
|
|
"os",
|
|
"sys",
|
|
"asyncio",
|
|
"json",
|
|
"logging",
|
|
"shutil",
|
|
"threading",
|
|
"time",
|
|
"typing",
|
|
"dataclasses",
|
|
"hashlib",
|
|
"re",
|
|
"datetime",
|
|
"fnmatch",
|
|
"random",
|
|
"collections",
|
|
"difflib",
|
|
"unittest",
|
|
"tempfile",
|
|
"abc",
|
|
"inspect",
|
|
"pathlib",
|
|
"sqlite3",
|
|
"base64",
|
|
"io",
|
|
"math",
|
|
"copy",
|
|
"itertools",
|
|
"functools",
|
|
"enum",
|
|
"struct",
|
|
"pickle",
|
|
"uuid",
|
|
"warnings",
|
|
"weakref",
|
|
"glob",
|
|
"cgi",
|
|
"contextlib",
|
|
"bisect",
|
|
"heapq",
|
|
"array",
|
|
"queue",
|
|
"concurrent",
|
|
"subprocess",
|
|
"socket",
|
|
"select",
|
|
"signal",
|
|
"errno",
|
|
"traceback",
|
|
"linecache",
|
|
"tokenize",
|
|
"token",
|
|
"keyword",
|
|
"operator",
|
|
"types",
|
|
"mimetypes",
|
|
"urllib",
|
|
"http",
|
|
"ftplib",
|
|
"poplib",
|
|
"imaplib",
|
|
"smtplib",
|
|
"telnetlib",
|
|
"nntplib",
|
|
"xml",
|
|
"html",
|
|
"csv",
|
|
"configparser",
|
|
"netrc",
|
|
"xdrlib",
|
|
"plistlib",
|
|
"hmac",
|
|
"secrets",
|
|
"multiprocessing",
|
|
"platform",
|
|
"resource",
|
|
"syslog",
|
|
"posix",
|
|
"pwd",
|
|
"grp",
|
|
"termios",
|
|
"tty",
|
|
"pty",
|
|
"fcntl",
|
|
"pipes",
|
|
"getpass",
|
|
"curses",
|
|
"dbm",
|
|
"gdbm",
|
|
# ComfyUI related
|
|
"folder_paths",
|
|
"server",
|
|
"comfy",
|
|
"nodes",
|
|
"node_helpers",
|
|
"execution",
|
|
"app",
|
|
"api",
|
|
"main",
|
|
"cli_args",
|
|
# ComfyUI prerequisite packages
|
|
"torch",
|
|
"torchsde",
|
|
"torchvision",
|
|
"torchaudio",
|
|
"numpy",
|
|
"PIL",
|
|
"aiohttp",
|
|
"yarl",
|
|
"yaml",
|
|
"tqdm",
|
|
"requests",
|
|
"psutil",
|
|
"einops",
|
|
"transformers",
|
|
"tokenizers",
|
|
"sentencepiece",
|
|
"safetensors",
|
|
"alembic",
|
|
"sqlalchemy",
|
|
"av",
|
|
"kornia",
|
|
"spandrel",
|
|
"pydantic",
|
|
"git",
|
|
"github",
|
|
"nio",
|
|
"huggingface_hub",
|
|
"typer",
|
|
"rich",
|
|
"typing_extensions",
|
|
"toml",
|
|
"uv",
|
|
"chardet",
|
|
# Allowed external dependencies in requirements.txt
|
|
"imagehash",
|
|
"scipy",
|
|
"pywt",
|
|
"zipfile",
|
|
}
|
|
|
|
|
|
class DependencyChecker(ASTChecker):
|
|
def __init__(self, file_path: str) -> None:
|
|
super().__init__(file_path)
|
|
self.found_imports: set[str] = set()
|
|
|
|
def get_target_files(self) -> list[str]:
|
|
# Logic to find all .py files in py/ and root __init__.py
|
|
targets = []
|
|
for root, _, files in os.walk("py"):
|
|
if any(x in root for x in ["node_modules", ".git", "__pycache__"]):
|
|
continue
|
|
for file in files:
|
|
if file.endswith(".py"):
|
|
targets.append(os.path.join(root, file))
|
|
|
|
if os.path.exists("__init__.py"):
|
|
targets.append("__init__.py")
|
|
return targets
|
|
|
|
def visit_Import(self, node: ast.Import) -> None:
|
|
for n in node.names:
|
|
self._check_import(n.name, node.lineno)
|
|
self.generic_visit(node)
|
|
|
|
def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
|
|
if node.level == 0 and node.module:
|
|
module: str = node.module
|
|
self._check_import(module, node.lineno)
|
|
self.generic_visit(node)
|
|
|
|
def _check_import(self, module_name: str, lineno: int) -> None:
|
|
imp = module_name.split(".")[0]
|
|
# Allow self-referencing modules
|
|
if imp == "py" or imp == "":
|
|
return
|
|
|
|
if imp not in ALLOWED_PACKAGES and not self._is_standard_library(imp):
|
|
# Dynamic check for standard library
|
|
try:
|
|
import importlib.util
|
|
|
|
spec = importlib.util.find_spec(imp)
|
|
if spec and (spec.origin == "built-in" or "lib/python" in (spec.origin or "")):
|
|
return
|
|
except (ImportError, ValueError):
|
|
pass
|
|
|
|
self.add_error(lineno, f"Unallowed dependency found: '{imp}'")
|
|
|
|
def _is_standard_library(self, module_name: str) -> bool:
|
|
return module_name in sys.builtin_module_names
|