Initial commit

This commit is contained in:
pacnpal
2024-11-19 01:21:25 -05:00
commit ec39193430
6 changed files with 530 additions and 0 deletions

5
pip_add/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
from .cli import main
from .utils import find_requirements, add_to_requirements
__version__ = "0.1.0"
__all__ = ['main', 'find_requirements', 'add_to_requirements']

238
pip_add/cli.py Normal file
View File

@@ -0,0 +1,238 @@
import os
import sys
import subprocess
import pkg_resources
import argparse
from pkg_resources import working_set
def find_requirements():
"""Find requirements.txt in current directory or create it"""
req_file = 'requirements.txt'
if not os.path.exists(req_file):
with open(req_file, 'w') as f:
f.write('# Python dependencies\n')
return req_file
def get_package_dependencies(package_name):
"""Get all dependencies of a package"""
try:
dist = pkg_resources.get_distribution(package_name)
deps = {package_name: dist.version} # Start with the main package
for req in dist.requires():
try:
dep_dist = pkg_resources.get_distribution(req.key)
deps[req.key] = dep_dist.version
except pkg_resources.DistributionNotFound:
continue
return deps
except Exception as e:
print(f"Warning: Could not fetch dependencies: {str(e)}", file=sys.stderr)
return {package_name: pkg_resources.get_distribution(package_name).version}
def find_dependent_packages(package_name, excluding_package):
"""
Find all installed packages that depend on the given package,
excluding the package we're removing and its dependencies
"""
dependents = {}
for dist in working_set:
# Skip the package we're removing and its dependencies
if dist.key == excluding_package or dist.key == package_name:
continue
for req in dist.requires():
if req.key == package_name:
dependents[package_name] = dependents.get(package_name, [])
dependents[package_name].append(dist.key)
return dependents
def analyze_dependencies(package_name):
"""Analyze which dependencies can be removed and which are still needed"""
try:
# Get all dependencies of the package we're removing
all_deps = get_package_dependencies(package_name)
safe_to_remove = set()
kept_deps = {}
# First, mark the main package for removal
main_pkg_version = all_deps.pop(package_name)
safe_to_remove.add((package_name, main_pkg_version))
# Then check each dependency
for dep, version in all_deps.items():
# Check if anything besides the package we're removing needs this dependency
dependents = find_dependent_packages(dep, package_name)
if not dependents.get(dep):
safe_to_remove.add((dep, version))
else:
kept_deps[dep] = dependents[dep]
return safe_to_remove, kept_deps
except Exception as e:
print(f"Warning: Could not analyze dependencies: {str(e)}", file=sys.stderr)
version = pkg_resources.get_distribution(package_name).version
return {(package_name, version)}, {}
def remove_from_requirements(packages, req_file):
"""Remove packages from requirements.txt"""
try:
with open(req_file, 'r') as f:
requirements = f.readlines()
except FileNotFoundError:
return set()
# Keep track of removed packages
removed = set()
# Filter out the packages while keeping comments
new_reqs = []
for line in requirements:
if line.strip() and not line.strip().startswith('#'):
# Check if this line is one of our packages
pkg = line.split('>=')[0].split('==')[0].strip()
if pkg not in {name for name, _ in packages}:
new_reqs.append(line)
else:
removed.add(pkg)
else:
new_reqs.append(line)
# Write back the filtered requirements
with open(req_file, 'w') as f:
f.writelines(new_reqs)
return removed
def add_to_requirements(packages_dict, req_file, exact=False):
"""Add or update packages in requirements.txt"""
try:
with open(req_file, 'r') as f:
requirements = f.readlines()
except FileNotFoundError:
requirements = []
# Keep comments and empty lines
filtered_reqs = [line for line in requirements if line.strip() and not line.strip().startswith('#')]
comments = [line for line in requirements if not line.strip() or line.strip().startswith('#')]
# Convert existing requirements to dict for easy updating
existing_pkgs = {}
for line in filtered_reqs:
if '>=' in line or '==' in line:
name = line.split('>=')[0].split('==')[0].strip()
existing_pkgs[name] = line
# Update or add new packages
for package, version in packages_dict.items():
operator = '==' if exact else '>='
requirement_line = f"{package}{operator}{version}\n"
existing_pkgs[package] = requirement_line
# Combine comments and updated requirements
final_requirements = comments + [existing_pkgs[pkg] for pkg in sorted(existing_pkgs.keys())]
with open(req_file, 'w') as f:
f.writelines(final_requirements)
def install_package(package_name):
"""Install a package using pip"""
print(f"Installing {package_name}...")
subprocess.check_call(
[sys.executable, '-m', 'pip', 'install', package_name],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL
)
def uninstall_packages(packages):
"""Uninstall multiple packages silently"""
print("Removing packages...")
for package, version in packages:
try:
subprocess.check_call(
[sys.executable, '-m', 'pip', 'uninstall', '-y', package],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL
)
except subprocess.CalledProcessError as e:
print(f"Warning: Error uninstalling {package}: {str(e)}", file=sys.stderr)
def main():
parser = argparse.ArgumentParser(description='Install/remove packages and manage requirements.txt')
parser.add_argument('package', help='Package to install or remove')
parser.add_argument('-d', '--with-dependencies', action='store_true',
help='Include dependencies when installing or removing')
parser.add_argument('-e', '--exact', action='store_true',
help='Use == instead of >= for version specification')
parser.add_argument('-r', '--remove', action='store_true',
help='Remove package(s) and their entries from requirements.txt')
args = parser.parse_args()
package = args.package
req_file = find_requirements()
try:
if args.remove:
if args.with_dependencies:
# Analyze dependencies
to_remove, kept_deps = analyze_dependencies(package)
# Remove from requirements.txt and uninstall
removed = remove_from_requirements(to_remove, req_file)
uninstall_packages(to_remove)
# Report results
if removed:
print("\n✓ Successfully uninstalled:")
for pkg, version in sorted(to_remove):
print(f" - {pkg} ({version})")
if kept_deps:
print("\n Dependencies kept (required by other packages):")
for dep, dependents in sorted(kept_deps.items()):
print(f" - {dep} (needed by: {', '.join(sorted(dependents))})")
print(f"\n✓ Updated {req_file}")
else:
print(f"\n✓ Successfully uninstalled {package} (no entries found in {req_file})")
else:
# Remove single package
version = pkg_resources.get_distribution(package).version
uninstall_packages({(package, version)})
removed = remove_from_requirements({(package, version)}, req_file)
if removed:
print(f"\n✓ Successfully uninstalled {package} ({version})")
print(f"✓ Updated {req_file}")
else:
print(f"\n✓ Successfully uninstalled {package} ({version})")
print(f" Note: Package was not found in {req_file}")
else:
# Install package
install_package(package)
# Get package info and update requirements
if args.with_dependencies:
packages = get_package_dependencies(package)
print("\n✓ Successfully installed:")
for pkg, version in sorted(packages.items()):
print(f" - {pkg} ({version})")
else:
version = pkg_resources.get_distribution(package).version
packages = {package: version}
print(f"\n✓ Successfully installed {package} ({version})")
# Update requirements.txt
add_to_requirements(packages, req_file, args.exact)
print(f"✓ Updated {req_file}")
except subprocess.CalledProcessError as e:
print(f"Error {'removing' if args.remove else 'installing'} {package}: {str(e)}", file=sys.stderr)
sys.exit(1)
except Exception as e:
print(f"Error: {str(e)}", file=sys.stderr)
sys.exit(1)
if __name__ == '__main__':
main()

83
pip_add/utils.py Normal file
View File

@@ -0,0 +1,83 @@
import os
from pathlib import Path
from typing import Optional, List, Tuple
def find_requirements(start_dir: Optional[str] = None) -> Tuple[str, List[str]]:
"""
Find requirements files by searching up through parent directories.
Returns tuple of (chosen_file, all_found_files).
"""
requirements_files = [
'requirements.txt',
'requirements/base.txt',
'requirements/dev.txt',
'requirements/development.txt',
'requirements/prod.txt',
'requirements/production.txt'
]
if start_dir is None:
start_dir = os.getcwd()
current_dir = Path(start_dir).resolve()
found_files = []
# Search up through parent directories
while True:
for req_file in requirements_files:
req_path = current_dir / req_file
if req_path.exists():
found_files.append(str(req_path))
# Stop if we found files or reached root
if found_files or current_dir.parent == current_dir:
break
current_dir = current_dir.parent
# If no files found, default to requirements.txt in current directory
if not found_files:
default_file = os.path.join(start_dir, 'requirements.txt')
# Create the requirements directory if needed
os.makedirs(os.path.dirname(default_file), exist_ok=True)
# Create an empty requirements file
if not os.path.exists(default_file):
with open(default_file, 'w') as f:
f.write('# Python dependencies\n')
return default_file, [default_file]
# Prefer requirements.txt in the closest directory
for f in found_files:
if os.path.basename(f) == 'requirements.txt':
return f, found_files
# Otherwise take the first found file
return found_files[0], found_files
def add_to_requirements(package_name: str, version: str, requirements_file: str) -> None:
"""Add or update package in requirements file"""
requirement_line = f"{package_name}=={version}\n"
try:
with open(requirements_file, 'r') as f:
requirements = f.readlines()
except FileNotFoundError:
requirements = []
# Check if package exists and update it
package_exists = False
for i, line in enumerate(requirements):
if line.strip() and package_name == line.split('==')[0].strip():
requirements[i] = requirement_line
package_exists = True
break
# Add package if it doesn't exist
if not package_exists:
requirements.append(requirement_line)
# Create directory if it doesn't exist
os.makedirs(os.path.dirname(os.path.abspath(requirements_file)), exist_ok=True)
with open(requirements_file, 'w') as f:
f.writelines(requirements)