166 lines
5.3 KiB
Python
166 lines
5.3 KiB
Python
from typing import Dict, List, Tuple, Any, Callable, TypeVar, Generic, NamedTuple
|
|
from packaging.utils import canonicalize_name, canonicalize_version
|
|
|
|
T = TypeVar("T") # For the package data type
|
|
K = TypeVar("K") # For the key type
|
|
V = TypeVar("V") # For the value type
|
|
|
|
|
|
class PackageLocation(NamedTuple):
|
|
"""
|
|
Composite key representing package name and location.
|
|
"""
|
|
|
|
name: str
|
|
location: str
|
|
|
|
|
|
class EnvironmentDiffTracker(Generic[T, K, V]):
|
|
"""
|
|
Generic utility class to track changes in environment states before and
|
|
after operations. Can be used with any environment management system
|
|
(pip, npm, apt, docker, etc.).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
key_extractor: Callable[[T], K],
|
|
value_extractor: Callable[[T], V],
|
|
) -> None:
|
|
"""
|
|
Initialize a new environment diff tracker.
|
|
|
|
Args:
|
|
key_extractor: Function to extract the item identifier from an entry
|
|
value_extractor: Function to extract the version or other value to
|
|
compare
|
|
normalize_key: Optional function to normalize keys
|
|
(e.g., make lowercase)
|
|
"""
|
|
self._key_extractor = key_extractor
|
|
self._value_extractor = value_extractor
|
|
self._before_items: Dict[K, V] = {}
|
|
self._after_items: Dict[K, V] = {}
|
|
|
|
def set_before_state(self, items_data: List[T]) -> None:
|
|
"""
|
|
Set the before-operation environment state.
|
|
|
|
Args:
|
|
items_data: List of items in the format specific to the environment
|
|
"""
|
|
self._before_items = self._normalize_items_data(items_data)
|
|
|
|
def set_after_state(self, items_data: List[T]) -> None:
|
|
"""
|
|
Set the after-operation environment state.
|
|
|
|
Args:
|
|
items_data: List of items in the format specific to the environment
|
|
"""
|
|
self._after_items = self._normalize_items_data(items_data)
|
|
|
|
def get_diff(self) -> Tuple[Dict[K, V], Dict[K, V], Dict[K, Tuple[V, V]]]:
|
|
"""
|
|
Compute the difference between before and after environment states.
|
|
|
|
Returns:
|
|
Tuple containing:
|
|
- Dictionary of added items {key: value}
|
|
- Dictionary of removed items {key: value}
|
|
- Dictionary of updated items {key: (old_value, new_value)}
|
|
"""
|
|
before_keys = set(self._before_items.keys())
|
|
after_keys = set(self._after_items.keys())
|
|
|
|
# Find added and removed items
|
|
added_keys = after_keys - before_keys
|
|
removed_keys = before_keys - after_keys
|
|
|
|
# Find updated items (same key, different value)
|
|
common_keys = before_keys & after_keys
|
|
updated_keys = {
|
|
key: (self._before_items[key], self._after_items[key])
|
|
for key in common_keys
|
|
if self._before_items[key] != self._after_items[key]
|
|
}
|
|
|
|
# Create result dictionaries
|
|
added = {key: self._after_items[key] for key in added_keys}
|
|
removed = {key: self._before_items[key] for key in removed_keys}
|
|
updated = {key: updated_keys[key] for key in updated_keys}
|
|
|
|
return added, removed, updated
|
|
|
|
def _normalize_items_data(self, items_data: List[T]) -> Dict[K, V]:
|
|
"""
|
|
Normalize items data into a standardized dictionary format.
|
|
|
|
Args:
|
|
items_data: List of item data entries
|
|
|
|
Returns:
|
|
Dict mapping normalized item keys to their values
|
|
"""
|
|
result = {}
|
|
|
|
for item_info in items_data:
|
|
try:
|
|
key = self._key_extractor(item_info)
|
|
value = self._value_extractor(item_info)
|
|
result[key] = value
|
|
except (KeyError, TypeError, AttributeError):
|
|
# Skip entries that don't have the expected structure
|
|
continue
|
|
return result
|
|
|
|
|
|
class PipEnvironmentDiffTracker(
|
|
EnvironmentDiffTracker[Dict[str, Any], PackageLocation, str]
|
|
):
|
|
"""
|
|
Specialized diff tracker for pip package environments.
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
key_extractor=self._pip_key_extractor,
|
|
value_extractor=self._pip_value_extractor,
|
|
)
|
|
|
|
# TODO: handle errors in value extraction
|
|
|
|
def _pip_key_extractor(self, pkg: Dict[str, Any]) -> PackageLocation:
|
|
return PackageLocation(
|
|
name=canonicalize_name(pkg.get("name", "")),
|
|
location=pkg.get("location", ""),
|
|
)
|
|
|
|
def _pip_value_extractor(self, pkg: Dict[str, Any]) -> str:
|
|
return canonicalize_version(pkg.get("version", ""), strip_trailing_zero=False)
|
|
|
|
|
|
class NpmEnvironmentDiffTracker(
|
|
EnvironmentDiffTracker[Dict[str, Any], PackageLocation, str]
|
|
):
|
|
"""
|
|
Specialized diff tracker for npm package environments.
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
key_extractor=self._npm_key_extractor,
|
|
value_extractor=self._npm_value_extractor,
|
|
)
|
|
|
|
# TODO: handle errors in value extraction
|
|
|
|
def _npm_key_extractor(self, pkg: Dict[str, Any]) -> PackageLocation:
|
|
return PackageLocation(
|
|
name=pkg.get("name", ""),
|
|
location=pkg.get("location", ""),
|
|
)
|
|
|
|
def _npm_value_extractor(self, pkg: Dict[str, Any]) -> str:
|
|
return pkg.get("version", "")
|