mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-10 06:10:50 +08:00
118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
import importlib.resources
|
|
from fsspec.spec import AbstractFileSystem
|
|
from fsspec.registry import register_implementation
|
|
|
|
|
|
class PkgResourcesFileSystem(AbstractFileSystem):
|
|
"""
|
|
An fsspec filesystem for reading Python package resources.
|
|
|
|
Paths are expected in the format:
|
|
pkg://<package.name>/path/to/resource.txt
|
|
"""
|
|
|
|
protocol = "pkg"
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._traversables = {}
|
|
|
|
def _get_traversable(self, package_name):
|
|
"""Get or cache the root Traversable for a package."""
|
|
if package_name not in self._traversables:
|
|
try:
|
|
# Get the Traversable object for the root of the package
|
|
self._traversables[package_name] = importlib.resources.files(package_name)
|
|
except ModuleNotFoundError as e:
|
|
raise FileNotFoundError(f"Package '{package_name}' not found.") from e
|
|
return self._traversables[package_name]
|
|
|
|
def _resolve_path(self, path):
|
|
"""Split a pkg:// path into package name and resource path."""
|
|
# Remove protocol and leading slashes
|
|
path_no_proto = self._strip_protocol(path).lstrip('/')
|
|
|
|
if not path_no_proto:
|
|
raise ValueError("Path must include a package name.")
|
|
|
|
parts = path_no_proto.split('/', 1)
|
|
package_name = parts[0]
|
|
|
|
resource_path = parts[1] if len(parts) > 1 else ""
|
|
|
|
root = self._get_traversable(package_name)
|
|
|
|
# Resolve the final resource Traversable
|
|
if resource_path:
|
|
resource = root.joinpath(resource_path)
|
|
else:
|
|
resource = root
|
|
|
|
return resource
|
|
|
|
def _open(
|
|
self,
|
|
path,
|
|
mode="rb",
|
|
**kwargs,
|
|
):
|
|
"""Open a file for reading."""
|
|
if "w" in mode or "a" in mode or "x" in mode:
|
|
raise NotImplementedError("Only read mode is supported.")
|
|
|
|
try:
|
|
resource = self._resolve_path(path)
|
|
if not resource.is_file():
|
|
raise FileNotFoundError(f"Path is not a file: {path}")
|
|
return resource.open("rb")
|
|
except (ModuleNotFoundError, FileNotFoundError):
|
|
raise FileNotFoundError(f"Resource not found: {path}")
|
|
except Exception as e:
|
|
raise IOError(f"Failed to open resource {path}: {e}") from e
|
|
|
|
def ls(self, path, detail=False, **kwargs):
|
|
"""List contents of a package directory."""
|
|
try:
|
|
resource = self._resolve_path(path)
|
|
if not resource.is_dir():
|
|
# If it's a file, 'ls' should return info on that file
|
|
return [self.info(path)] if detail else [path]
|
|
|
|
items = []
|
|
for item in resource.iterdir():
|
|
item_path = f"{path.rstrip('/')}/{item.name}"
|
|
if detail:
|
|
items.append(self.info(item_path))
|
|
else:
|
|
items.append(item_path)
|
|
return items
|
|
except (ModuleNotFoundError, FileNotFoundError):
|
|
raise FileNotFoundError(f"Resource path not found: {path}")
|
|
|
|
def info(self, path, **kwargs):
|
|
"""Get info about a resource."""
|
|
try:
|
|
resource = self._resolve_path(path)
|
|
resource_type = "directory" if resource.is_dir() else "file"
|
|
|
|
size = None
|
|
if resource_type == 'file':
|
|
# This is inefficient but demonstrates the principle
|
|
try:
|
|
with resource.open('rb') as f:
|
|
size = len(f.read())
|
|
except Exception:
|
|
size = None # Could fail for some reason
|
|
|
|
return {
|
|
"name": path,
|
|
"type": resource_type,
|
|
"size": size,
|
|
}
|
|
except (ModuleNotFoundError, FileNotFoundError):
|
|
raise FileNotFoundError(f"Resource not found: {path}")
|
|
|
|
|
|
# Register the filesystem with fsspec
|
|
register_implementation(PkgResourcesFileSystem.protocol, PkgResourcesFileSystem)
|