Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,12 @@
|
||||
from .analyze.is_from_package import is_from_package
|
||||
from .file_structure_representation import Directory
|
||||
from .glob_group import GlobGroup
|
||||
from .importer import (
|
||||
Importer,
|
||||
ObjMismatchError,
|
||||
ObjNotFoundError,
|
||||
OrderedImporter,
|
||||
sys_importer,
|
||||
)
|
||||
from .package_exporter import EmptyMatchError, PackageExporter, PackagingError
|
||||
from .package_importer import PackageImporter
|
||||
@@ -0,0 +1,173 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections import deque
|
||||
|
||||
|
||||
class DiGraph:
|
||||
"""Really simple unweighted directed graph data structure to track dependencies.
|
||||
|
||||
The API is pretty much the same as networkx so if you add something just
|
||||
copy their API.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Dict of node -> dict of arbitrary attributes
|
||||
self._node = {}
|
||||
# Nested dict of node -> successor node -> nothing.
|
||||
# (didn't implement edge data)
|
||||
self._succ = {}
|
||||
# Nested dict of node -> predecessor node -> nothing.
|
||||
self._pred = {}
|
||||
|
||||
# Keep track of the order in which nodes are added to
|
||||
# the graph.
|
||||
self._node_order = {}
|
||||
self._insertion_idx = 0
|
||||
|
||||
def add_node(self, n, **kwargs):
|
||||
"""Add a node to the graph.
|
||||
|
||||
Args:
|
||||
n: the node. Can we any object that is a valid dict key.
|
||||
**kwargs: any attributes you want to attach to the node.
|
||||
"""
|
||||
if n not in self._node:
|
||||
self._node[n] = kwargs
|
||||
self._succ[n] = {}
|
||||
self._pred[n] = {}
|
||||
self._node_order[n] = self._insertion_idx
|
||||
self._insertion_idx += 1
|
||||
else:
|
||||
self._node[n].update(kwargs)
|
||||
|
||||
def add_edge(self, u, v):
|
||||
"""Add an edge to graph between nodes ``u`` and ``v``
|
||||
|
||||
``u`` and ``v`` will be created if they do not already exist.
|
||||
"""
|
||||
# add nodes
|
||||
self.add_node(u)
|
||||
self.add_node(v)
|
||||
|
||||
# add the edge
|
||||
self._succ[u][v] = True
|
||||
self._pred[v][u] = True
|
||||
|
||||
def successors(self, n):
|
||||
"""Returns an iterator over successor nodes of n."""
|
||||
try:
|
||||
return iter(self._succ[n])
|
||||
except KeyError as e:
|
||||
raise ValueError(f"The node {n} is not in the digraph.") from e
|
||||
|
||||
def predecessors(self, n):
|
||||
"""Returns an iterator over predecessors nodes of n."""
|
||||
try:
|
||||
return iter(self._pred[n])
|
||||
except KeyError as e:
|
||||
raise ValueError(f"The node {n} is not in the digraph.") from e
|
||||
|
||||
@property
|
||||
def edges(self):
|
||||
"""Returns an iterator over all edges (u, v) in the graph"""
|
||||
for n, successors in self._succ.items():
|
||||
for succ in successors:
|
||||
yield n, succ
|
||||
|
||||
@property
|
||||
def nodes(self):
|
||||
"""Returns a dictionary of all nodes to their attributes."""
|
||||
return self._node
|
||||
|
||||
def __iter__(self):
|
||||
"""Iterate over the nodes."""
|
||||
return iter(self._node)
|
||||
|
||||
def __contains__(self, n):
|
||||
"""Returns True if ``n`` is a node in the graph, False otherwise."""
|
||||
try:
|
||||
return n in self._node
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
def forward_transitive_closure(self, src: str) -> set[str]:
|
||||
"""Returns a set of nodes that are reachable from src"""
|
||||
|
||||
result = set(src)
|
||||
working_set = deque(src)
|
||||
while len(working_set) > 0:
|
||||
cur = working_set.popleft()
|
||||
for n in self.successors(cur):
|
||||
if n not in result:
|
||||
result.add(n)
|
||||
working_set.append(n)
|
||||
return result
|
||||
|
||||
def backward_transitive_closure(self, src: str) -> set[str]:
|
||||
"""Returns a set of nodes that are reachable from src in reverse direction"""
|
||||
|
||||
result = set(src)
|
||||
working_set = deque(src)
|
||||
while len(working_set) > 0:
|
||||
cur = working_set.popleft()
|
||||
for n in self.predecessors(cur):
|
||||
if n not in result:
|
||||
result.add(n)
|
||||
working_set.append(n)
|
||||
return result
|
||||
|
||||
def all_paths(self, src: str, dst: str):
|
||||
"""Returns a subgraph rooted at src that shows all the paths to dst."""
|
||||
|
||||
result_graph = DiGraph()
|
||||
# First compute forward transitive closure of src (all things reachable from src).
|
||||
forward_reachable_from_src = self.forward_transitive_closure(src)
|
||||
|
||||
if dst not in forward_reachable_from_src:
|
||||
return result_graph
|
||||
|
||||
# Second walk the reverse dependencies of dst, adding each node to
|
||||
# the output graph iff it is also present in forward_reachable_from_src.
|
||||
# we don't use backward_transitive_closures for optimization purposes
|
||||
working_set = deque(dst)
|
||||
while len(working_set) > 0:
|
||||
cur = working_set.popleft()
|
||||
for n in self.predecessors(cur):
|
||||
if n in forward_reachable_from_src:
|
||||
result_graph.add_edge(n, cur)
|
||||
# only explore further if its reachable from src
|
||||
working_set.append(n)
|
||||
|
||||
return result_graph.to_dot()
|
||||
|
||||
def first_path(self, dst: str) -> list[str]:
|
||||
"""Returns a list of nodes that show the first path that resulted in dst being added to the graph."""
|
||||
path = []
|
||||
|
||||
while dst:
|
||||
path.append(dst)
|
||||
candidates = self._pred[dst].keys()
|
||||
dst, min_idx = "", None
|
||||
for candidate in candidates:
|
||||
idx = self._node_order.get(candidate, None)
|
||||
if idx is None:
|
||||
break
|
||||
if min_idx is None or idx < min_idx:
|
||||
min_idx = idx
|
||||
dst = candidate
|
||||
|
||||
return list(reversed(path))
|
||||
|
||||
def to_dot(self) -> str:
|
||||
"""Returns the dot representation of the graph.
|
||||
|
||||
Returns:
|
||||
A dot representation of the graph.
|
||||
"""
|
||||
edges = "\n".join(f'"{f}" -> "{t}";' for f, t in self.edges)
|
||||
return f"""\
|
||||
digraph G {{
|
||||
rankdir = LR;
|
||||
node [shape=box];
|
||||
{edges}
|
||||
}}
|
||||
"""
|
||||
@@ -0,0 +1,66 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import os.path
|
||||
from glob import glob
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
from torch.types import Storage
|
||||
|
||||
|
||||
__serialization_id_record_name__ = ".data/serialization_id"
|
||||
|
||||
|
||||
# because get_storage_from_record returns a tensor!?
|
||||
class _HasStorage:
|
||||
def __init__(self, storage):
|
||||
self._storage = storage
|
||||
|
||||
def storage(self):
|
||||
return self._storage
|
||||
|
||||
|
||||
class DirectoryReader:
|
||||
"""
|
||||
Class to allow PackageImporter to operate on unzipped packages. Methods
|
||||
copy the behavior of the internal PyTorchFileReader class (which is used for
|
||||
accessing packages in all other cases).
|
||||
|
||||
N.B.: ScriptObjects are not depickleable or accessible via this DirectoryReader
|
||||
class due to ScriptObjects requiring an actual PyTorchFileReader instance.
|
||||
"""
|
||||
|
||||
def __init__(self, directory):
|
||||
self.directory = directory
|
||||
|
||||
def get_record(self, name):
|
||||
filename = f"{self.directory}/{name}"
|
||||
with open(filename, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
def get_storage_from_record(self, name, numel, dtype):
|
||||
filename = f"{self.directory}/{name}"
|
||||
nbytes = torch._utils._element_size(dtype) * numel
|
||||
storage = cast(Storage, torch.UntypedStorage)
|
||||
return _HasStorage(storage.from_file(filename=filename, nbytes=nbytes))
|
||||
|
||||
def has_record(self, path):
|
||||
full_path = os.path.join(self.directory, path)
|
||||
return os.path.isfile(full_path)
|
||||
|
||||
def get_all_records(
|
||||
self,
|
||||
):
|
||||
files = [
|
||||
filename[len(self.directory) + 1 :]
|
||||
for filename in glob(f"{self.directory}/**", recursive=True)
|
||||
if not os.path.isdir(filename)
|
||||
]
|
||||
return files
|
||||
|
||||
def serialization_id(
|
||||
self,
|
||||
):
|
||||
if self.has_record(__serialization_id_record_name__):
|
||||
return self.get_record(__serialization_id_record_name__)
|
||||
else:
|
||||
return ""
|
||||
@@ -0,0 +1,95 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import _warnings
|
||||
import os.path
|
||||
|
||||
|
||||
# note: implementations
|
||||
# copied from cpython's import code
|
||||
|
||||
|
||||
# _zip_searchorder defines how we search for a module in the Zip
|
||||
# archive: we first search for a package __init__, then for
|
||||
# non-package .pyc, and .py entries. The .pyc entries
|
||||
# are swapped by initzipimport() if we run in optimized mode. Also,
|
||||
# '/' is replaced by path_sep there.
|
||||
|
||||
_zip_searchorder = (
|
||||
("/__init__.py", True),
|
||||
(".py", False),
|
||||
)
|
||||
|
||||
|
||||
# Replace any occurrences of '\r\n?' in the input string with '\n'.
|
||||
# This converts DOS and Mac line endings to Unix line endings.
|
||||
def _normalize_line_endings(source):
|
||||
source = source.replace(b"\r\n", b"\n")
|
||||
source = source.replace(b"\r", b"\n")
|
||||
return source
|
||||
|
||||
|
||||
def _resolve_name(name, package, level):
|
||||
"""Resolve a relative module name to an absolute one."""
|
||||
bits = package.rsplit(".", level - 1)
|
||||
if len(bits) < level:
|
||||
raise ValueError("attempted relative import beyond top-level package")
|
||||
base = bits[0]
|
||||
return f"{base}.{name}" if name else base
|
||||
|
||||
|
||||
def _sanity_check(name, package, level):
|
||||
"""Verify arguments are "sane"."""
|
||||
if not isinstance(name, str):
|
||||
raise TypeError(f"module name must be str, not {type(name)}")
|
||||
if level < 0:
|
||||
raise ValueError("level must be >= 0")
|
||||
if level > 0:
|
||||
if not isinstance(package, str):
|
||||
raise TypeError("__package__ not set to a string")
|
||||
elif not package:
|
||||
raise ImportError("attempted relative import with no known parent package")
|
||||
if not name and level == 0:
|
||||
raise ValueError("Empty module name")
|
||||
|
||||
|
||||
def _calc___package__(globals):
|
||||
"""Calculate what __package__ should be.
|
||||
|
||||
__package__ is not guaranteed to be defined or could be set to None
|
||||
to represent that its proper value is unknown.
|
||||
|
||||
"""
|
||||
package = globals.get("__package__")
|
||||
spec = globals.get("__spec__")
|
||||
if package is not None:
|
||||
if spec is not None and package != spec.parent:
|
||||
_warnings.warn( # noqa: G010
|
||||
f"__package__ != __spec__.parent ({package!r} != {spec.parent!r})", # noqa: G004
|
||||
ImportWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
return package
|
||||
elif spec is not None:
|
||||
return spec.parent
|
||||
else:
|
||||
_warnings.warn( # noqa: G010
|
||||
"can't resolve package from __spec__ or __package__, "
|
||||
"falling back on __name__ and __path__",
|
||||
ImportWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
package = globals["__name__"]
|
||||
if "__path__" not in globals:
|
||||
package = package.rpartition(".")[0]
|
||||
return package
|
||||
|
||||
|
||||
def _normalize_path(path):
|
||||
"""Normalize a path by ensuring it is a string.
|
||||
|
||||
If the resulting string contains path separators, an exception is raised.
|
||||
"""
|
||||
parent, file_name = os.path.split(path)
|
||||
if parent:
|
||||
raise ValueError(f"{path!r} must be only a file name")
|
||||
else:
|
||||
return file_name
|
||||
@@ -0,0 +1,66 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Import mangling.
|
||||
See mangling.md for details.
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
|
||||
_mangle_index = 0
|
||||
|
||||
|
||||
class PackageMangler:
|
||||
"""
|
||||
Used on import, to ensure that all modules imported have a shared mangle parent.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
global _mangle_index
|
||||
self._mangle_index = _mangle_index
|
||||
# Increment the global index
|
||||
_mangle_index += 1
|
||||
# Angle brackets are used so that there is almost no chance of
|
||||
# confusing this module for a real module. Plus, it is Python's
|
||||
# preferred way of denoting special modules.
|
||||
self._mangle_parent = f"<torch_package_{self._mangle_index}>"
|
||||
|
||||
def mangle(self, name) -> str:
|
||||
if len(name) == 0:
|
||||
raise AssertionError("name must not be empty")
|
||||
return self._mangle_parent + "." + name
|
||||
|
||||
def demangle(self, mangled: str) -> str:
|
||||
"""
|
||||
Note: This only demangles names that were mangled by this specific
|
||||
PackageMangler. It will pass through names created by a different
|
||||
PackageMangler instance.
|
||||
"""
|
||||
if mangled.startswith(self._mangle_parent + "."):
|
||||
return mangled.partition(".")[2]
|
||||
|
||||
# wasn't a mangled name
|
||||
return mangled
|
||||
|
||||
def parent_name(self):
|
||||
return self._mangle_parent
|
||||
|
||||
|
||||
def is_mangled(name: str) -> bool:
|
||||
return bool(re.match(r"<torch_package_\d+>", name))
|
||||
|
||||
|
||||
def demangle(name: str) -> str:
|
||||
"""
|
||||
Note: Unlike PackageMangler.demangle, this version works on any
|
||||
mangled name, irrespective of which PackageMangler created it.
|
||||
"""
|
||||
if is_mangled(name):
|
||||
_first, sep, last = name.partition(".")
|
||||
# If there is only a base mangle prefix, e.g. '<torch_package_0>',
|
||||
# then return an empty string.
|
||||
return last if len(sep) != 0 else ""
|
||||
return name
|
||||
|
||||
|
||||
def get_mangle_prefix(name: str) -> str:
|
||||
return name.partition(".")[0] if is_mangled(name) else name
|
||||
@@ -0,0 +1,123 @@
|
||||
# mypy: allow-untyped-defs
|
||||
_magic_methods = [
|
||||
"__subclasscheck__",
|
||||
"__hex__",
|
||||
"__rmul__",
|
||||
"__float__",
|
||||
"__idiv__",
|
||||
"__setattr__",
|
||||
"__div__",
|
||||
"__invert__",
|
||||
"__nonzero__",
|
||||
"__rshift__",
|
||||
"__eq__",
|
||||
"__pos__",
|
||||
"__round__",
|
||||
"__rand__",
|
||||
"__or__",
|
||||
"__complex__",
|
||||
"__divmod__",
|
||||
"__len__",
|
||||
"__reversed__",
|
||||
"__copy__",
|
||||
"__reduce__",
|
||||
"__deepcopy__",
|
||||
"__rdivmod__",
|
||||
"__rrshift__",
|
||||
"__ifloordiv__",
|
||||
"__hash__",
|
||||
"__iand__",
|
||||
"__xor__",
|
||||
"__isub__",
|
||||
"__oct__",
|
||||
"__ceil__",
|
||||
"__imod__",
|
||||
"__add__",
|
||||
"__truediv__",
|
||||
"__unicode__",
|
||||
"__le__",
|
||||
"__delitem__",
|
||||
"__sizeof__",
|
||||
"__sub__",
|
||||
"__ne__",
|
||||
"__pow__",
|
||||
"__bytes__",
|
||||
"__mul__",
|
||||
"__itruediv__",
|
||||
"__bool__",
|
||||
"__iter__",
|
||||
"__abs__",
|
||||
"__gt__",
|
||||
"__iadd__",
|
||||
"__enter__",
|
||||
"__floordiv__",
|
||||
"__call__",
|
||||
"__neg__",
|
||||
"__and__",
|
||||
"__ixor__",
|
||||
"__getitem__",
|
||||
"__exit__",
|
||||
"__cmp__",
|
||||
"__getstate__",
|
||||
"__index__",
|
||||
"__contains__",
|
||||
"__floor__",
|
||||
"__lt__",
|
||||
"__getattr__",
|
||||
"__mod__",
|
||||
"__trunc__",
|
||||
"__delattr__",
|
||||
"__instancecheck__",
|
||||
"__setitem__",
|
||||
"__ipow__",
|
||||
"__ilshift__",
|
||||
"__long__",
|
||||
"__irshift__",
|
||||
"__imul__",
|
||||
"__lshift__",
|
||||
"__dir__",
|
||||
"__ge__",
|
||||
"__int__",
|
||||
"__ior__",
|
||||
]
|
||||
|
||||
|
||||
class MockedObject:
|
||||
_name: str
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
# _suppress_err is set by us in the mocked module impl, so that we can
|
||||
# construct instances of MockedObject to hand out to people looking up
|
||||
# module attributes.
|
||||
|
||||
# Any other attempt to construct a MockedObject instance (say, in the
|
||||
# unpickling process) should give an error.
|
||||
if not kwargs.get("_suppress_err"):
|
||||
raise NotImplementedError(
|
||||
f"Object '{cls._name}' was mocked out during packaging "
|
||||
f"but it is being used in '__new__'. If this error is "
|
||||
"happening during 'load_pickle', please ensure that your "
|
||||
"pickled object doesn't contain any mocked objects."
|
||||
)
|
||||
# Otherwise, this is just a regular object creation
|
||||
# (e.g. `x = MockedObject("foo")`), so pass it through normally.
|
||||
return super().__new__(cls)
|
||||
|
||||
def __init__(self, name: str, _suppress_err: bool):
|
||||
self.__dict__["_name"] = name
|
||||
|
||||
def __repr__(self):
|
||||
return f"MockedObject({self._name})"
|
||||
|
||||
|
||||
def install_method(method_name):
|
||||
def _not_implemented(self, *args, **kwargs):
|
||||
raise NotImplementedError(
|
||||
f"Object '{self._name}' was mocked out during packaging but it is being used in {method_name}"
|
||||
)
|
||||
|
||||
setattr(MockedObject, method_name, _not_implemented)
|
||||
|
||||
|
||||
for method_name in _magic_methods:
|
||||
install_method(method_name)
|
||||
@@ -0,0 +1,144 @@
|
||||
import sys
|
||||
from pickle import (
|
||||
_compat_pickle, # pyrefly: ignore [missing-module-attribute]
|
||||
_extension_registry, # pyrefly: ignore [missing-module-attribute]
|
||||
_getattribute, # pyrefly: ignore [missing-module-attribute]
|
||||
_Pickler,
|
||||
EXT1,
|
||||
EXT2,
|
||||
EXT4,
|
||||
GLOBAL,
|
||||
PicklingError,
|
||||
STACK_GLOBAL,
|
||||
)
|
||||
from struct import pack
|
||||
from types import FunctionType
|
||||
|
||||
from .importer import Importer, ObjMismatchError, ObjNotFoundError, sys_importer
|
||||
|
||||
|
||||
class _PyTorchLegacyPickler(_Pickler):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._persistent_id = None
|
||||
|
||||
def persistent_id(self, obj):
|
||||
if self._persistent_id is None:
|
||||
return super().persistent_id(obj)
|
||||
return self._persistent_id(obj)
|
||||
|
||||
|
||||
class PackagePickler(_PyTorchLegacyPickler):
|
||||
"""Package-aware pickler.
|
||||
|
||||
This behaves the same as a normal pickler, except it uses an `Importer`
|
||||
to find objects and modules to save.
|
||||
"""
|
||||
|
||||
def __init__(self, importer: Importer, *args, **kwargs):
|
||||
self.importer = importer
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
# Make sure the dispatch table copied from _Pickler is up-to-date.
|
||||
# Previous issues have been encountered where a library (e.g. dill)
|
||||
# mutate _Pickler.dispatch, PackagePickler makes a copy when this lib
|
||||
# is imported, then the offending library removes its dispatch entries,
|
||||
# leaving PackagePickler with a stale dispatch table that may cause
|
||||
# unwanted behavior.
|
||||
self.dispatch = _Pickler.dispatch.copy() # type: ignore[misc]
|
||||
self.dispatch[FunctionType] = PackagePickler.save_global # type: ignore[assignment]
|
||||
|
||||
def save_global(self, obj, name=None):
|
||||
# ruff: noqa: F841
|
||||
# unfortunately the pickler code is factored in a way that
|
||||
# forces us to copy/paste this function. The only change is marked
|
||||
# CHANGED below.
|
||||
write = self.write # type: ignore[attr-defined]
|
||||
memo = self.memo # type: ignore[attr-defined]
|
||||
|
||||
# CHANGED: import module from module environment instead of __import__
|
||||
try:
|
||||
module_name, name = self.importer.get_name(obj, name)
|
||||
except (ObjNotFoundError, ObjMismatchError) as err:
|
||||
raise PicklingError(f"Can't pickle {obj}: {str(err)}") from err
|
||||
|
||||
module = self.importer.import_module(module_name)
|
||||
if sys.version_info >= (3, 14):
|
||||
# pickle._getattribute signature changes in 3.14
|
||||
# to take iterable and return just the object (not tuple)
|
||||
# We need to get the parent object that contains the attribute
|
||||
name_parts = name.split(".")
|
||||
if "<locals>" in name_parts:
|
||||
raise PicklingError(f"Can't pickle local object {obj!r}")
|
||||
if len(name_parts) == 1:
|
||||
parent = module
|
||||
else:
|
||||
parent = _getattribute(module, name_parts[:-1])
|
||||
else:
|
||||
_, parent = _getattribute(module, name)
|
||||
# END CHANGED
|
||||
|
||||
if self.proto >= 2: # type: ignore[attr-defined]
|
||||
code = _extension_registry.get((module_name, name))
|
||||
if code:
|
||||
if code <= 0:
|
||||
raise AssertionError(
|
||||
f"expected positive extension code, got {code}"
|
||||
)
|
||||
if code <= 0xFF:
|
||||
write(EXT1 + pack("<B", code))
|
||||
elif code <= 0xFFFF:
|
||||
write(EXT2 + pack("<H", code))
|
||||
else:
|
||||
write(EXT4 + pack("<i", code))
|
||||
return
|
||||
lastname = name.rpartition(".")[2]
|
||||
if parent is module:
|
||||
name = lastname
|
||||
# Non-ASCII identifiers are supported only with protocols >= 3.
|
||||
if self.proto >= 4: # type: ignore[attr-defined]
|
||||
self.save(module_name) # type: ignore[attr-defined]
|
||||
self.save(name) # type: ignore[attr-defined]
|
||||
write(STACK_GLOBAL)
|
||||
elif parent is not module:
|
||||
self.save_reduce(getattr, (parent, lastname)) # type: ignore[attr-defined]
|
||||
elif self.proto >= 3: # type: ignore[attr-defined]
|
||||
write(
|
||||
GLOBAL
|
||||
+ bytes(module_name, "utf-8")
|
||||
+ b"\n"
|
||||
+ bytes(name, "utf-8")
|
||||
+ b"\n"
|
||||
)
|
||||
else:
|
||||
if self.fix_imports: # type: ignore[attr-defined]
|
||||
r_name_mapping = _compat_pickle.REVERSE_NAME_MAPPING
|
||||
r_import_mapping = _compat_pickle.REVERSE_IMPORT_MAPPING
|
||||
if (module_name, name) in r_name_mapping:
|
||||
module_name, name = r_name_mapping[(module_name, name)]
|
||||
elif module_name in r_import_mapping:
|
||||
module_name = r_import_mapping[module_name]
|
||||
try:
|
||||
write(
|
||||
GLOBAL
|
||||
+ bytes(module_name, "ascii")
|
||||
+ b"\n"
|
||||
+ bytes(name, "ascii")
|
||||
+ b"\n"
|
||||
)
|
||||
except UnicodeEncodeError as exc:
|
||||
raise PicklingError(
|
||||
f"can't pickle global identifier '{module}.{name}' using "
|
||||
f"pickle protocol {self.proto:d}" # type: ignore[attr-defined]
|
||||
) from exc
|
||||
|
||||
self.memoize(obj) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def create_pickler(data_buf, importer, protocol=4):
|
||||
if importer is sys_importer:
|
||||
# if we are using the normal import library system, then
|
||||
# we can use the C implementation of pickle which is faster
|
||||
return _PyTorchLegacyPickler(data_buf, protocol=protocol)
|
||||
else:
|
||||
return PackagePickler(importer, data_buf, protocol=protocol)
|
||||
@@ -0,0 +1,27 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import _compat_pickle
|
||||
import pickle
|
||||
|
||||
from .importer import Importer
|
||||
|
||||
|
||||
class PackageUnpickler(pickle._Unpickler): # type: ignore[name-defined]
|
||||
"""Package-aware unpickler.
|
||||
|
||||
This behaves the same as a normal unpickler, except it uses `importer` to
|
||||
find any global names that it encounters while unpickling.
|
||||
"""
|
||||
|
||||
def __init__(self, importer: Importer, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._importer = importer
|
||||
|
||||
def find_class(self, module, name):
|
||||
# Subclasses may override this.
|
||||
if self.proto < 3 and self.fix_imports: # type: ignore[attr-defined]
|
||||
if (module, name) in _compat_pickle.NAME_MAPPING:
|
||||
module, name = _compat_pickle.NAME_MAPPING[(module, name)]
|
||||
elif module in _compat_pickle.IMPORT_MAPPING:
|
||||
module = _compat_pickle.IMPORT_MAPPING[module]
|
||||
mod = self._importer.import_module(module)
|
||||
return getattr(mod, name)
|
||||
@@ -0,0 +1,20 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""List of Python standard library modules.
|
||||
|
||||
Sadly, there is no reliable way to tell whether a module is part of the
|
||||
standard library except by comparing to a canonical list.
|
||||
|
||||
This is taken from https://github.com/PyCQA/isort/tree/develop/isort/stdlibs,
|
||||
which itself is sourced from the Python documentation.
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
|
||||
def is_stdlib_module(module: str) -> bool:
|
||||
base_module = module.partition(".")[0]
|
||||
return base_module in _get_stdlib_modules()
|
||||
|
||||
|
||||
def _get_stdlib_modules():
|
||||
return sys.stdlib_module_names
|
||||
@@ -0,0 +1,2 @@
|
||||
from .find_first_use_of_broken_modules import find_first_use_of_broken_modules
|
||||
from .trace_dependencies import trace_dependencies
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
from torch.package.package_exporter import PackagingError
|
||||
|
||||
|
||||
__all__ = ["find_first_use_of_broken_modules"]
|
||||
|
||||
|
||||
def find_first_use_of_broken_modules(exc: PackagingError) -> dict[str, list[str]]:
|
||||
"""
|
||||
Find all broken modules in a PackagingError, and for each one, return the
|
||||
dependency path in which the module was first encountered.
|
||||
|
||||
E.g. broken module m.n.o was added to a dependency graph while processing a.b.c,
|
||||
then re-encountered while processing d.e.f. This method would return
|
||||
{'m.n.o': ['a', 'b', 'c']}
|
||||
|
||||
Args:
|
||||
exc: a PackagingError
|
||||
|
||||
Returns: A dict from broken module names to lists of module names in the path.
|
||||
"""
|
||||
|
||||
if not isinstance(exc, PackagingError):
|
||||
raise AssertionError(
|
||||
f"exception must be a PackagingError, got {type(exc).__name__}"
|
||||
)
|
||||
uses = {}
|
||||
broken_module_names = [
|
||||
m for m, attr in exc.dependency_graph.nodes.items() if attr.get("error", False)
|
||||
]
|
||||
for module_name in broken_module_names:
|
||||
path = exc.dependency_graph.first_path(module_name)
|
||||
uses[module_name] = path
|
||||
return uses
|
||||
@@ -0,0 +1,16 @@
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
|
||||
from .._mangling import is_mangled
|
||||
|
||||
|
||||
def is_from_package(obj: Any) -> bool:
|
||||
"""
|
||||
Return whether an object was loaded from a package.
|
||||
|
||||
Note: packaged objects from externed modules will return ``False``.
|
||||
"""
|
||||
if type(obj) is ModuleType:
|
||||
return is_mangled(obj.__name__)
|
||||
else:
|
||||
return is_mangled(type(obj).__module__)
|
||||
@@ -0,0 +1,65 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import sys
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import Any
|
||||
|
||||
|
||||
__all__ = ["trace_dependencies"]
|
||||
|
||||
|
||||
def trace_dependencies(
|
||||
callable: Callable[[Any], Any], inputs: Iterable[tuple[Any, ...]]
|
||||
) -> list[str]:
|
||||
"""Trace the execution of a callable in order to determine which modules it uses.
|
||||
|
||||
Args:
|
||||
callable: The callable to execute and trace.
|
||||
inputs: The input to use during tracing. The modules used by 'callable' when invoked by each set of inputs
|
||||
are union-ed to determine all modules used by the callable for the purpooses of packaging.
|
||||
|
||||
Returns: A list of the names of all modules used during callable execution.
|
||||
"""
|
||||
modules_used = set()
|
||||
|
||||
def record_used_modules(frame, event, arg):
|
||||
# If the event being profiled is not a Python function
|
||||
# call, there is nothing to do.
|
||||
if event != "call":
|
||||
return
|
||||
|
||||
# This is the name of the function that was called.
|
||||
name = frame.f_code.co_name
|
||||
module = None
|
||||
|
||||
# Try to determine the name of the module that the function
|
||||
# is in:
|
||||
# 1) Check the global namespace of the frame.
|
||||
# 2) Check the local namespace of the frame.
|
||||
# 3) To handle class instance method calls, check
|
||||
# the attribute named 'name' of the object
|
||||
# in the local namespace corresponding to "self".
|
||||
if name in frame.f_globals:
|
||||
module = frame.f_globals[name].__module__
|
||||
elif name in frame.f_locals:
|
||||
module = frame.f_locals[name].__module__
|
||||
elif "self" in frame.f_locals:
|
||||
method = getattr(frame.f_locals["self"], name, None)
|
||||
module = method.__module__ if method else None
|
||||
|
||||
# If a module was found, add it to the set of used modules.
|
||||
if module:
|
||||
modules_used.add(module)
|
||||
|
||||
try:
|
||||
# Attach record_used_modules as the profiler function.
|
||||
sys.setprofile(record_used_modules)
|
||||
|
||||
# Execute the callable with all inputs.
|
||||
for inp in inputs:
|
||||
callable(*inp)
|
||||
|
||||
finally:
|
||||
# Detach the profiler function.
|
||||
sys.setprofile(None)
|
||||
|
||||
return list(modules_used)
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
from .glob_group import GlobGroup, GlobPattern
|
||||
|
||||
|
||||
__all__ = ["Directory"]
|
||||
|
||||
|
||||
class Directory:
|
||||
"""A file structure representation. Organized as Directory nodes that have lists of
|
||||
their Directory children. Directories for a package are created by calling
|
||||
:meth:`PackageImporter.file_structure`."""
|
||||
|
||||
def __init__(self, name: str, is_dir: bool):
|
||||
self.name = name
|
||||
self.is_dir = is_dir
|
||||
self.children: dict[str, Directory] = {}
|
||||
|
||||
def _get_dir(self, dirs: list[str]) -> "Directory":
|
||||
"""Builds path of Directories if not yet built and returns last directory
|
||||
in list.
|
||||
|
||||
Args:
|
||||
dirs (List[str]): List of directory names that are treated like a path.
|
||||
|
||||
Returns:
|
||||
:class:`Directory`: The last Directory specified in the dirs list.
|
||||
"""
|
||||
if len(dirs) == 0:
|
||||
return self
|
||||
dir_name = dirs[0]
|
||||
if dir_name not in self.children:
|
||||
self.children[dir_name] = Directory(dir_name, True)
|
||||
return self.children[dir_name]._get_dir(dirs[1:])
|
||||
|
||||
def _add_file(self, file_path: str):
|
||||
"""Adds a file to a Directory.
|
||||
|
||||
Args:
|
||||
file_path (str): Path of file to add. Last element is added as a file while
|
||||
other paths items are added as directories.
|
||||
"""
|
||||
*dirs, file = file_path.split("/")
|
||||
dir = self._get_dir(dirs)
|
||||
dir.children[file] = Directory(file, False)
|
||||
|
||||
def has_file(self, filename: str) -> bool:
|
||||
"""Checks if a file is present in a :class:`Directory`.
|
||||
|
||||
Args:
|
||||
filename (str): Path of file to search for.
|
||||
Returns:
|
||||
bool: If a :class:`Directory` contains the specified file.
|
||||
"""
|
||||
lineage = filename.split("/", maxsplit=1)
|
||||
child = lineage[0]
|
||||
grandchildren = lineage[1] if len(lineage) > 1 else None
|
||||
if child in self.children:
|
||||
if grandchildren is None:
|
||||
return True
|
||||
else:
|
||||
return self.children[child].has_file(grandchildren)
|
||||
return False
|
||||
|
||||
def __str__(self):
|
||||
str_list: list[str] = []
|
||||
self._stringify_tree(str_list)
|
||||
return "".join(str_list)
|
||||
|
||||
def _stringify_tree(
|
||||
self,
|
||||
str_list: list[str],
|
||||
preamble: str = "",
|
||||
dir_ptr: str = "\u2500\u2500\u2500 ",
|
||||
):
|
||||
"""Recursive method to generate print-friendly version of a Directory."""
|
||||
space = " "
|
||||
branch = "\u2502 "
|
||||
tee = "\u251c\u2500\u2500 "
|
||||
last = "\u2514\u2500\u2500 "
|
||||
|
||||
# add this directory's representation
|
||||
str_list.append(f"{preamble}{dir_ptr}{self.name}\n")
|
||||
|
||||
# add directory's children representations
|
||||
if dir_ptr == tee:
|
||||
preamble = preamble + branch
|
||||
else:
|
||||
preamble = preamble + space
|
||||
|
||||
file_keys: list[str] = []
|
||||
dir_keys: list[str] = []
|
||||
for key, val in self.children.items():
|
||||
if val.is_dir:
|
||||
dir_keys.append(key)
|
||||
else:
|
||||
file_keys.append(key)
|
||||
|
||||
for index, key in enumerate(sorted(dir_keys)):
|
||||
if (index == len(dir_keys) - 1) and len(file_keys) == 0:
|
||||
self.children[key]._stringify_tree(str_list, preamble, last)
|
||||
else:
|
||||
self.children[key]._stringify_tree(str_list, preamble, tee)
|
||||
for index, file in enumerate(sorted(file_keys)):
|
||||
pointer = last if (index == len(file_keys) - 1) else tee
|
||||
str_list.append(f"{preamble}{pointer}{file}\n")
|
||||
|
||||
|
||||
def _create_directory_from_file_list(
|
||||
filename: str,
|
||||
file_list: list[str],
|
||||
include: "GlobPattern" = "**",
|
||||
exclude: "GlobPattern" = (),
|
||||
) -> Directory:
|
||||
"""Return a :class:`Directory` file structure representation created from a list of files.
|
||||
|
||||
Args:
|
||||
filename (str): The name given to the top-level directory that will be the
|
||||
relative root for all file paths found in the file_list.
|
||||
|
||||
file_list (List[str]): List of files to add to the top-level directory.
|
||||
|
||||
include (list[str] | str): An optional pattern that limits what is included from the file_list to
|
||||
files whose name matches the pattern.
|
||||
|
||||
exclude (list[str] | str): An optional pattern that excludes files whose name match the pattern.
|
||||
|
||||
Returns:
|
||||
:class:`Directory`: a :class:`Directory` file structure representation created from a list of files.
|
||||
"""
|
||||
glob_pattern = GlobGroup(include, exclude=exclude, separator="/")
|
||||
|
||||
top_dir = Directory(filename, True)
|
||||
for file in file_list:
|
||||
if glob_pattern.matches(file):
|
||||
top_dir._add_file(file)
|
||||
return top_dir
|
||||
@@ -0,0 +1,95 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import ast
|
||||
|
||||
from ._importlib import _resolve_name
|
||||
|
||||
|
||||
class _ExtractModuleReferences(ast.NodeVisitor):
|
||||
"""
|
||||
Extract the list of global variables a block of code will read and write
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def run(cls, src: str, package: str) -> list[tuple[str, str | None]]:
|
||||
visitor = cls(package)
|
||||
tree = ast.parse(src)
|
||||
visitor.visit(tree)
|
||||
return list(visitor.references.keys())
|
||||
|
||||
def __init__(self, package):
|
||||
super().__init__()
|
||||
self.package = package
|
||||
self.references = {}
|
||||
|
||||
def _absmodule(self, module_name: str, level: int) -> str:
|
||||
if level > 0:
|
||||
return _resolve_name(module_name, self.package, level)
|
||||
return module_name
|
||||
|
||||
def visit_Import(self, node):
|
||||
for alias in node.names:
|
||||
self.references[(alias.name, None)] = True
|
||||
|
||||
def visit_ImportFrom(self, node):
|
||||
name = self._absmodule(node.module, 0 if node.level is None else node.level)
|
||||
for alias in node.names:
|
||||
# from my_package import foo
|
||||
# foo may be a module, so we have to add it to the list of
|
||||
# potential references, if import of it fails, we will ignore it
|
||||
if alias.name != "*":
|
||||
self.references[(name, alias.name)] = True
|
||||
else:
|
||||
self.references[(name, None)] = True
|
||||
|
||||
def _grab_node_int(self, node):
|
||||
return node.value
|
||||
|
||||
def _grab_node_str(self, node):
|
||||
return node.value
|
||||
|
||||
def visit_Call(self, node):
|
||||
# __import__ calls aren't routed to the visit_Import/From nodes
|
||||
if hasattr(node.func, "id") and node.func.id == "__import__":
|
||||
try:
|
||||
name = self._grab_node_str(node.args[0])
|
||||
fromlist: list[str] = []
|
||||
level = 0
|
||||
if len(node.args) > 3:
|
||||
fromlist.extend(self._grab_node_str(v) for v in node.args[3].elts)
|
||||
elif hasattr(node, "keywords"):
|
||||
for keyword in node.keywords:
|
||||
if keyword.arg == "fromlist":
|
||||
fromlist.extend(
|
||||
self._grab_node_str(v) for v in keyword.value.elts
|
||||
)
|
||||
if len(node.args) > 4:
|
||||
level = self._grab_node_int(node.args[4])
|
||||
elif hasattr(node, "keywords"):
|
||||
for keyword in node.keywords:
|
||||
if keyword.arg == "level":
|
||||
level = self._grab_node_int(keyword.value)
|
||||
if fromlist == []:
|
||||
# the top-level package (the name up till the first dot) is returned
|
||||
# when the fromlist argument is empty in normal import system,
|
||||
# we need to include top level package to match this behavior and last
|
||||
# level package to capture the intended dependency of user
|
||||
self.references[(name, None)] = True
|
||||
top_name = name.rsplit(".", maxsplit=1)[0]
|
||||
if top_name != name:
|
||||
top_name = self._absmodule(top_name, level)
|
||||
self.references[(top_name, None)] = True
|
||||
else:
|
||||
name = self._absmodule(name, level)
|
||||
for alias in fromlist:
|
||||
# fromlist args may be submodules, so we have to add the fromlist args
|
||||
# to the list of potential references. If import of an arg fails we
|
||||
# will ignore it, similar to visit_ImportFrom
|
||||
if alias != "*":
|
||||
self.references[(name, alias)] = True
|
||||
else:
|
||||
self.references[(name, None)] = True
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
find_files_source_depends_on = _ExtractModuleReferences.run
|
||||
@@ -0,0 +1,86 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Iterable
|
||||
|
||||
|
||||
GlobPattern = str | Iterable[str]
|
||||
|
||||
|
||||
class GlobGroup:
|
||||
"""A set of patterns that candidate strings will be matched against.
|
||||
|
||||
A candidate is composed of a list of segments separated by ``separator``, e.g. "foo.bar.baz".
|
||||
|
||||
A pattern contains one or more segments. Segments can be:
|
||||
- A literal string (e.g. "foo"), which matches exactly.
|
||||
- A string containing a wildcard (e.g. "torch*", or "foo*baz*"). The wildcard matches
|
||||
any string, including the empty string.
|
||||
- A double wildcard ("**"). This matches against zero or more complete segments.
|
||||
|
||||
Examples:
|
||||
``torch.**``: matches ``torch`` and all its submodules, e.g. ``torch.nn`` and ``torch.nn.functional``.
|
||||
``torch.*``: matches ``torch.nn`` or ``torch.functional``, but not ``torch.nn.functional``.
|
||||
``torch*.**``: matches ``torch``, ``torchvision``, and all their submodules.
|
||||
|
||||
A candidates will match the ``GlobGroup`` if it matches any of the ``include`` patterns and
|
||||
none of the ``exclude`` patterns.
|
||||
|
||||
Args:
|
||||
include (str | Iterable[str]): A string or list of strings,
|
||||
each representing a pattern to be matched against. A candidate
|
||||
will match if it matches *any* include pattern
|
||||
exclude (str | Iterable[str]): A string or list of strings,
|
||||
each representing a pattern to be matched against. A candidate
|
||||
will be excluded from matching if it matches *any* exclude pattern.
|
||||
separator (str): A string that delimits segments in candidates and
|
||||
patterns. By default this is "." which corresponds to how modules are
|
||||
named in Python. Another common value for this is "/", which is
|
||||
the Unix path separator.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, include: GlobPattern, *, exclude: GlobPattern = (), separator: str = "."
|
||||
):
|
||||
self._dbg = f"GlobGroup(include={include}, exclude={exclude})"
|
||||
self.include = GlobGroup._glob_list(include, separator)
|
||||
self.exclude = GlobGroup._glob_list(exclude, separator)
|
||||
self.separator = separator
|
||||
|
||||
def __str__(self):
|
||||
return self._dbg
|
||||
|
||||
def __repr__(self):
|
||||
return self._dbg
|
||||
|
||||
def matches(self, candidate: str) -> bool:
|
||||
candidate = self.separator + candidate
|
||||
return any(p.fullmatch(candidate) for p in self.include) and all(
|
||||
not p.fullmatch(candidate) for p in self.exclude
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _glob_list(elems: GlobPattern, separator: str = "."):
|
||||
if isinstance(elems, str):
|
||||
return [GlobGroup._glob_to_re(elems, separator)]
|
||||
else:
|
||||
return [GlobGroup._glob_to_re(e, separator) for e in elems]
|
||||
|
||||
@staticmethod
|
||||
def _glob_to_re(pattern: str, separator: str = "."):
|
||||
# to avoid corner cases for the first component, we prefix the candidate string
|
||||
# with '.' so `import torch` will regex against `.torch`, assuming '.' is the separator
|
||||
def component_to_re(component):
|
||||
if "**" in component:
|
||||
if component == "**":
|
||||
return "(" + re.escape(separator) + "[^" + separator + "]+)*"
|
||||
else:
|
||||
raise ValueError("** can only appear as an entire path segment")
|
||||
else:
|
||||
return re.escape(separator) + ("[^" + separator + "]*").join(
|
||||
re.escape(x) for x in component.split("*")
|
||||
)
|
||||
|
||||
result = "".join(component_to_re(c) for c in pattern.split(separator))
|
||||
return re.compile(result)
|
||||
@@ -0,0 +1,262 @@
|
||||
import importlib
|
||||
import logging
|
||||
import sys
|
||||
from abc import ABC, abstractmethod
|
||||
from pickle import (
|
||||
_getattribute, # pyrefly: ignore [missing-module-attribute]
|
||||
_Pickler,
|
||||
whichmodule as _pickle_whichmodule, # pyrefly: ignore [missing-module-attribute]
|
||||
)
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
|
||||
from ._mangling import demangle, get_mangle_prefix, is_mangled
|
||||
|
||||
|
||||
__all__ = ["ObjNotFoundError", "ObjMismatchError", "Importer", "OrderedImporter"]
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ObjNotFoundError(Exception):
|
||||
"""Raised when an importer cannot find an object by searching for its name."""
|
||||
|
||||
|
||||
class ObjMismatchError(Exception):
|
||||
"""Raised when an importer found a different object with the same name as the user-provided one."""
|
||||
|
||||
|
||||
class Importer(ABC):
|
||||
"""Represents an environment to import modules from.
|
||||
|
||||
By default, you can figure out what module an object belongs by checking
|
||||
__module__ and importing the result using __import__ or importlib.import_module.
|
||||
|
||||
torch.package introduces module importers other than the default one.
|
||||
Each PackageImporter introduces a new namespace. Potentially a single
|
||||
name (e.g. 'foo.bar') is present in multiple namespaces.
|
||||
|
||||
It supports two main operations:
|
||||
import_module: module_name -> module object
|
||||
get_name: object -> (parent module name, name of obj within module)
|
||||
|
||||
The guarantee is that following round-trip will succeed or throw an ObjNotFoundError/ObjMisMatchError.
|
||||
module_name, obj_name = env.get_name(obj)
|
||||
module = env.import_module(module_name)
|
||||
obj2 = getattr(module, obj_name)
|
||||
assert obj1 is obj2
|
||||
"""
|
||||
|
||||
modules: dict[str, ModuleType]
|
||||
|
||||
@abstractmethod
|
||||
def import_module(self, module_name: str) -> ModuleType:
|
||||
"""Import `module_name` from this environment.
|
||||
|
||||
The contract is the same as for importlib.import_module.
|
||||
"""
|
||||
|
||||
def get_name(self, obj: Any, name: str | None = None) -> tuple[str, str]:
|
||||
"""Given an object, return a name that can be used to retrieve the
|
||||
object from this environment.
|
||||
|
||||
Args:
|
||||
obj: An object to get the module-environment-relative name for.
|
||||
name: If set, use this name instead of looking up __name__ or __qualname__ on `obj`.
|
||||
This is only here to match how Pickler handles __reduce__ functions that return a string,
|
||||
don't use otherwise.
|
||||
Returns:
|
||||
A tuple (parent_module_name, attr_name) that can be used to retrieve `obj` from this environment.
|
||||
Use it like:
|
||||
mod = importer.import_module(parent_module_name)
|
||||
obj = getattr(mod, attr_name)
|
||||
|
||||
Raises:
|
||||
ObjNotFoundError: we couldn't retrieve `obj by name.
|
||||
ObjMisMatchError: we found a different object with the same name as `obj`.
|
||||
"""
|
||||
if name is None and obj and _Pickler.dispatch.get(type(obj)) is None:
|
||||
# Honor the string return variant of __reduce__, which will give us
|
||||
# a global name to search for in this environment.
|
||||
# TODO: I guess we should do copyreg too?
|
||||
reduce = getattr(obj, "__reduce__", None)
|
||||
if reduce is not None:
|
||||
try:
|
||||
rv = reduce()
|
||||
if isinstance(rv, str):
|
||||
name = rv
|
||||
except Exception:
|
||||
pass
|
||||
if name is None:
|
||||
name = getattr(obj, "__qualname__", None)
|
||||
if name is None:
|
||||
name = obj.__name__
|
||||
|
||||
orig_module_name = self.whichmodule(obj, name)
|
||||
# Demangle the module name before importing. If this obj came out of a
|
||||
# PackageImporter, `__module__` will be mangled. See mangling.md for
|
||||
# details.
|
||||
module_name = demangle(orig_module_name)
|
||||
|
||||
# Check that this name will indeed return the correct object
|
||||
try:
|
||||
module = self.import_module(module_name)
|
||||
if sys.version_info >= (3, 14):
|
||||
# pickle._getatribute signature changes in 3.14
|
||||
# to take iterable and return just one object
|
||||
obj2 = _getattribute(module, name.split("."))
|
||||
else:
|
||||
obj2, _ = _getattribute(module, name)
|
||||
except (ImportError, KeyError, AttributeError):
|
||||
raise ObjNotFoundError(
|
||||
f"{obj} was not found as {module_name}.{name}"
|
||||
) from None
|
||||
|
||||
if obj is obj2:
|
||||
return module_name, name
|
||||
|
||||
def get_obj_info(obj):
|
||||
if name is None:
|
||||
raise AssertionError("name must not be None")
|
||||
module_name = self.whichmodule(obj, name)
|
||||
is_mangled_ = is_mangled(module_name)
|
||||
location = (
|
||||
get_mangle_prefix(module_name)
|
||||
if is_mangled_
|
||||
else "the current Python environment"
|
||||
)
|
||||
importer_name = (
|
||||
f"the importer for {get_mangle_prefix(module_name)}"
|
||||
if is_mangled_
|
||||
else "'sys_importer'"
|
||||
)
|
||||
return module_name, location, importer_name
|
||||
|
||||
obj_module_name, obj_location, obj_importer_name = get_obj_info(obj)
|
||||
obj2_module_name, obj2_location, obj2_importer_name = get_obj_info(obj2)
|
||||
msg = (
|
||||
f"\n\nThe object provided is from '{obj_module_name}', "
|
||||
f"which is coming from {obj_location}."
|
||||
f"\nHowever, when we import '{obj2_module_name}', it's coming from {obj2_location}."
|
||||
"\nTo fix this, make sure this 'PackageExporter's importer lists "
|
||||
f"{obj_importer_name} before {obj2_importer_name}."
|
||||
)
|
||||
raise ObjMismatchError(msg)
|
||||
|
||||
def whichmodule(self, obj: Any, name: str) -> str:
|
||||
"""Find the module name an object belongs to.
|
||||
|
||||
This should be considered internal for end-users, but developers of
|
||||
an importer can override it to customize the behavior.
|
||||
|
||||
Taken from pickle.py, but modified to exclude the search into sys.modules
|
||||
"""
|
||||
module_name = getattr(obj, "__module__", None)
|
||||
if module_name is not None:
|
||||
return module_name
|
||||
|
||||
# Protect the iteration by using a list copy of self.modules against dynamic
|
||||
# modules that trigger imports of other modules upon calls to getattr.
|
||||
for module_name, module in self.modules.copy().items():
|
||||
if (
|
||||
module_name == "__main__"
|
||||
or module_name == "__mp_main__" # bpo-42406
|
||||
or module is None
|
||||
):
|
||||
continue
|
||||
try:
|
||||
if _getattribute(module, name)[0] is obj:
|
||||
return module_name
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
return "__main__"
|
||||
|
||||
|
||||
class _SysImporter(Importer):
|
||||
"""An importer that implements the default behavior of Python."""
|
||||
|
||||
def import_module(self, module_name: str):
|
||||
return importlib.import_module(module_name)
|
||||
|
||||
def whichmodule(self, obj: Any, name: str) -> str:
|
||||
# In Python 3.14+, pickle.whichmodule tries to import the module,
|
||||
# which fails for mangled package names like '<torch_package_0>'.
|
||||
# Check __module__ first before calling pickle.whichmodule.
|
||||
module_name = getattr(obj, "__module__", None)
|
||||
if module_name is not None:
|
||||
return module_name
|
||||
return _pickle_whichmodule(obj, name)
|
||||
|
||||
|
||||
sys_importer = _SysImporter()
|
||||
|
||||
|
||||
class OrderedImporter(Importer):
|
||||
"""A compound importer that takes a list of importers and tries them one at a time.
|
||||
|
||||
The first importer in the list that returns a result "wins".
|
||||
"""
|
||||
|
||||
def __init__(self, *args):
|
||||
self._importers: list[Importer] = list(args)
|
||||
|
||||
def _is_torchpackage_dummy(self, module):
|
||||
"""Returns true iff this module is an empty PackageNode in a torch.package.
|
||||
|
||||
If you intern `a.b` but never use `a` in your code, then `a` will be an
|
||||
empty module with no source. This can break cases where we are trying to
|
||||
re-package an object after adding a real dependency on `a`, since
|
||||
OrderedImportere will resolve `a` to the dummy package and stop there.
|
||||
|
||||
See: https://github.com/pytorch/pytorch/pull/71520#issuecomment-1029603769
|
||||
"""
|
||||
if not getattr(module, "__torch_package__", False):
|
||||
return False
|
||||
if not hasattr(module, "__path__"):
|
||||
return False
|
||||
if not hasattr(module, "__file__"):
|
||||
return True
|
||||
return module.__file__ is None
|
||||
|
||||
def get_name(self, obj: Any, name: str | None = None) -> tuple[str, str]:
|
||||
for importer in self._importers:
|
||||
try:
|
||||
return importer.get_name(obj, name)
|
||||
except (ObjNotFoundError, ObjMismatchError) as e:
|
||||
warning_message = (
|
||||
f"Tried to call get_name with obj {obj}, "
|
||||
f"and name {name} on {importer} and got {e}"
|
||||
)
|
||||
log.warning(warning_message)
|
||||
raise ObjNotFoundError(
|
||||
f"Could not find obj {obj} and name {name} in any of the importers {self._importers}"
|
||||
)
|
||||
|
||||
def import_module(self, module_name: str) -> ModuleType:
|
||||
last_err = None
|
||||
for importer in self._importers:
|
||||
if not isinstance(importer, Importer):
|
||||
raise TypeError(
|
||||
f"{importer} is not a Importer. "
|
||||
"All importers in OrderedImporter must inherit from Importer."
|
||||
)
|
||||
try:
|
||||
module = importer.import_module(module_name)
|
||||
if self._is_torchpackage_dummy(module):
|
||||
continue
|
||||
return module
|
||||
except ModuleNotFoundError as err:
|
||||
last_err = err
|
||||
|
||||
if last_err is not None:
|
||||
raise last_err
|
||||
else:
|
||||
raise ModuleNotFoundError(module_name)
|
||||
|
||||
def whichmodule(self, obj: Any, name: str) -> str:
|
||||
for importer in self._importers:
|
||||
module_name = importer.whichmodule(obj, name)
|
||||
if module_name != "__main__":
|
||||
return module_name
|
||||
|
||||
return "__main__"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,806 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import builtins
|
||||
import importlib
|
||||
import importlib.machinery
|
||||
import inspect
|
||||
import io
|
||||
import linecache
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
from collections.abc import Callable, Iterable
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, cast, TYPE_CHECKING
|
||||
from weakref import WeakValueDictionary
|
||||
|
||||
import torch
|
||||
from torch.serialization import _get_restore_location, _maybe_decode_ascii
|
||||
from torch.types import FileLike
|
||||
|
||||
from ._directory_reader import DirectoryReader
|
||||
from ._importlib import (
|
||||
_calc___package__,
|
||||
_normalize_line_endings,
|
||||
_normalize_path,
|
||||
_resolve_name,
|
||||
_sanity_check,
|
||||
)
|
||||
from ._mangling import demangle, PackageMangler
|
||||
from ._package_unpickler import PackageUnpickler
|
||||
from .file_structure_representation import _create_directory_from_file_list, Directory
|
||||
from .importer import Importer
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .glob_group import GlobPattern
|
||||
|
||||
__all__ = ["PackageImporter"]
|
||||
|
||||
|
||||
# This is a list of imports that are implicitly allowed even if they haven't
|
||||
# been marked as extern. This is to work around the fact that Torch implicitly
|
||||
# depends on numpy and package can't track it.
|
||||
# https://github.com/pytorch/multipy/issues/46 # codespell:ignore multipy
|
||||
IMPLICIT_IMPORT_ALLOWLIST: Iterable[str] = [
|
||||
"numpy",
|
||||
"numpy.core",
|
||||
"numpy.core._multiarray_umath",
|
||||
# FX GraphModule might depend on builtins module and users usually
|
||||
# don't extern builtins. Here we import it here by default.
|
||||
"builtins",
|
||||
]
|
||||
|
||||
|
||||
# Compatibility name mapping to facilitate upgrade of external modules.
|
||||
# The primary motivation is to enable Numpy upgrade that many modules
|
||||
# depend on. The latest release of Numpy removed `numpy.str` and
|
||||
# `numpy.bool` breaking unpickling for many modules.
|
||||
EXTERN_IMPORT_COMPAT_NAME_MAPPING: dict[str, dict[str, Any]] = {
|
||||
"numpy": {
|
||||
"str": str,
|
||||
"bool": bool,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class PackageImporter(Importer):
|
||||
"""Importers allow you to load code written to packages by :class:`PackageExporter`.
|
||||
Code is loaded in a hermetic way, using files from the package
|
||||
rather than the normal python import system. This allows
|
||||
for the packaging of PyTorch model code and data so that it can be run
|
||||
on a server or used in the future for transfer learning.
|
||||
|
||||
The importer for packages ensures that code in the module can only be loaded from
|
||||
within the package, except for modules explicitly listed as external during export.
|
||||
The file ``extern_modules`` in the zip archive lists all the modules that a package externally depends on.
|
||||
This prevents "implicit" dependencies where the package runs locally because it is importing
|
||||
a locally-installed package, but then fails when the package is copied to another machine.
|
||||
"""
|
||||
|
||||
"""The dictionary of already loaded modules from this package, equivalent to ``sys.modules`` but
|
||||
local to this importer.
|
||||
"""
|
||||
|
||||
modules: dict[str, types.ModuleType]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_or_buffer: FileLike | torch._C.PyTorchFileReader,
|
||||
module_allowed: Callable[[str], bool] = lambda module_name: True,
|
||||
):
|
||||
"""Open ``file_or_buffer`` for importing. This checks that the imported package only requires modules
|
||||
allowed by ``module_allowed``
|
||||
|
||||
Args:
|
||||
file_or_buffer: a file-like object (has to implement :meth:`read`, :meth:`readline`, :meth:`tell`, and :meth:`seek`),
|
||||
a string, or an ``os.PathLike`` object containing a filename.
|
||||
module_allowed (Callable[[str], bool], optional): A method to determine if a externally provided module
|
||||
should be allowed. Can be used to ensure packages loaded do not depend on modules that the server
|
||||
does not support. Defaults to allowing anything.
|
||||
|
||||
Raises:
|
||||
ImportError: If the package will use a disallowed module.
|
||||
"""
|
||||
torch._C._log_api_usage_once("torch.package.PackageImporter")
|
||||
|
||||
self.zip_reader: Any
|
||||
if isinstance(file_or_buffer, torch._C.PyTorchFileReader):
|
||||
self.filename = "<pytorch_file_reader>"
|
||||
self.zip_reader = file_or_buffer
|
||||
elif isinstance(file_or_buffer, (os.PathLike, str)):
|
||||
self.filename = os.fspath(file_or_buffer)
|
||||
if not os.path.isdir(self.filename):
|
||||
self.zip_reader = torch._C.PyTorchFileReader(self.filename)
|
||||
else:
|
||||
self.zip_reader = DirectoryReader(self.filename)
|
||||
else:
|
||||
self.filename = "<binary>"
|
||||
self.zip_reader = torch._C.PyTorchFileReader(file_or_buffer)
|
||||
|
||||
torch._C._log_api_usage_metadata(
|
||||
"torch.package.PackageImporter.metadata",
|
||||
{
|
||||
"serialization_id": self.zip_reader.serialization_id(),
|
||||
"file_name": self.filename,
|
||||
},
|
||||
)
|
||||
|
||||
self.root = _PackageNode(None)
|
||||
self.modules = {}
|
||||
self.extern_modules = self._read_extern()
|
||||
|
||||
for extern_module in self.extern_modules:
|
||||
if not module_allowed(extern_module):
|
||||
raise ImportError(
|
||||
f"package '{file_or_buffer}' needs the external module '{extern_module}' "
|
||||
f"but that module has been disallowed"
|
||||
)
|
||||
self._add_extern(extern_module)
|
||||
|
||||
for fname in self.zip_reader.get_all_records():
|
||||
self._add_file(fname)
|
||||
|
||||
self.patched_builtins = builtins.__dict__.copy()
|
||||
self.patched_builtins["__import__"] = self.__import__
|
||||
# Allow packaged modules to reference their PackageImporter
|
||||
self.modules["torch_package_importer"] = self # type: ignore[assignment]
|
||||
|
||||
self._mangler = PackageMangler()
|
||||
|
||||
# used for reduce deserializaiton
|
||||
self.storage_context: Any = None
|
||||
self.last_map_location = None
|
||||
|
||||
# used for torch.serialization._load
|
||||
self.Unpickler = lambda *args, **kwargs: PackageUnpickler(self, *args, **kwargs)
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
def import_module(self, name: str, package=None):
|
||||
"""Load a module from the package if it hasn't already been loaded, and then return
|
||||
the module. Modules are loaded locally
|
||||
to the importer and will appear in ``self.modules`` rather than ``sys.modules``.
|
||||
|
||||
Args:
|
||||
name (str): Fully qualified name of the module to load.
|
||||
package ([type], optional): Unused, but present to match the signature of importlib.import_module. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
types.ModuleType: The (possibly already) loaded module.
|
||||
"""
|
||||
# We should always be able to support importing modules from this package.
|
||||
# This is to support something like:
|
||||
# obj = importer.load_pickle(...)
|
||||
# importer.import_module(obj.__module__) <- this string will be mangled
|
||||
#
|
||||
# Note that _mangler.demangle will not demangle any module names
|
||||
# produced by a different PackageImporter instance.
|
||||
name = self._mangler.demangle(name)
|
||||
|
||||
return self._gcd_import(name)
|
||||
|
||||
def load_binary(self, package: str, resource: str) -> bytes:
|
||||
"""Load raw bytes.
|
||||
|
||||
Args:
|
||||
package (str): The name of module package (e.g. ``"my_package.my_subpackage"``).
|
||||
resource (str): The unique name for the resource.
|
||||
|
||||
Returns:
|
||||
bytes: The loaded data.
|
||||
"""
|
||||
|
||||
path = self._zipfile_path(package, resource)
|
||||
return self.zip_reader.get_record(path)
|
||||
|
||||
def load_text(
|
||||
self,
|
||||
package: str,
|
||||
resource: str,
|
||||
encoding: str = "utf-8",
|
||||
errors: str = "strict",
|
||||
) -> str:
|
||||
"""Load a string.
|
||||
|
||||
Args:
|
||||
package (str): The name of module package (e.g. ``"my_package.my_subpackage"``).
|
||||
resource (str): The unique name for the resource.
|
||||
encoding (str, optional): Passed to ``decode``. Defaults to ``'utf-8'``.
|
||||
errors (str, optional): Passed to ``decode``. Defaults to ``'strict'``.
|
||||
|
||||
Returns:
|
||||
str: The loaded text.
|
||||
"""
|
||||
data = self.load_binary(package, resource)
|
||||
return data.decode(encoding, errors)
|
||||
|
||||
def load_pickle(self, package: str, resource: str, map_location=None) -> Any:
|
||||
"""Unpickles the resource from the package, loading any modules that are needed to construct the objects
|
||||
using :meth:`import_module`.
|
||||
|
||||
Args:
|
||||
package (str): The name of module package (e.g. ``"my_package.my_subpackage"``).
|
||||
resource (str): The unique name for the resource.
|
||||
map_location: Passed to `torch.load` to determine how tensors are mapped to devices. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
Any: The unpickled object.
|
||||
"""
|
||||
pickle_file = self._zipfile_path(package, resource)
|
||||
restore_location = _get_restore_location(map_location)
|
||||
loaded_storages = {}
|
||||
loaded_reduces = {}
|
||||
storage_context = torch._C.DeserializationStorageContext()
|
||||
|
||||
def load_tensor(dtype, size, key, location, restore_location):
|
||||
name = f"{key}.storage"
|
||||
|
||||
if storage_context.has_storage(name):
|
||||
storage = storage_context.get_storage(name, dtype)._typed_storage()
|
||||
else:
|
||||
tensor = self.zip_reader.get_storage_from_record(
|
||||
".data/" + name, size, dtype
|
||||
)
|
||||
if isinstance(self.zip_reader, torch._C.PyTorchFileReader):
|
||||
storage_context.add_storage(name, tensor)
|
||||
storage = tensor._typed_storage()
|
||||
loaded_storages[key] = restore_location(storage, location)
|
||||
|
||||
def persistent_load(saved_id):
|
||||
if not isinstance(saved_id, tuple):
|
||||
raise AssertionError(
|
||||
f"saved_id must be a tuple, got {type(saved_id).__name__}"
|
||||
)
|
||||
typename = _maybe_decode_ascii(saved_id[0])
|
||||
data = saved_id[1:]
|
||||
|
||||
if typename == "storage":
|
||||
storage_type, key, location, size = data
|
||||
if storage_type is torch.UntypedStorage:
|
||||
dtype = torch.uint8
|
||||
else:
|
||||
dtype = storage_type.dtype
|
||||
|
||||
if key not in loaded_storages:
|
||||
load_tensor(
|
||||
dtype,
|
||||
size,
|
||||
key,
|
||||
_maybe_decode_ascii(location),
|
||||
restore_location,
|
||||
)
|
||||
storage = loaded_storages[key]
|
||||
# TODO: Once we decide to break serialization FC, we can
|
||||
# stop wrapping with TypedStorage
|
||||
return torch.storage.TypedStorage(
|
||||
wrap_storage=storage._untyped_storage, dtype=dtype, _internal=True
|
||||
)
|
||||
elif typename == "reduce_package":
|
||||
# to fix BC breaking change, objects on this load path
|
||||
# will be loaded multiple times erroneously
|
||||
if len(data) == 2:
|
||||
func, args = data
|
||||
return func(self, *args)
|
||||
reduce_id, func, args = data
|
||||
if reduce_id not in loaded_reduces:
|
||||
loaded_reduces[reduce_id] = func(self, *args)
|
||||
return loaded_reduces[reduce_id]
|
||||
else:
|
||||
f"Unknown typename for persistent_load, expected 'storage' or 'reduce_package' but got '{typename}'"
|
||||
|
||||
# Load the data (which may in turn use `persistent_load` to load tensors)
|
||||
data_file = io.BytesIO(self.zip_reader.get_record(pickle_file))
|
||||
unpickler = self.Unpickler(data_file)
|
||||
unpickler.persistent_load = persistent_load # type: ignore[assignment]
|
||||
|
||||
@contextmanager
|
||||
def set_deserialization_context():
|
||||
# to let reduce_package access deserializaiton context
|
||||
self.storage_context = storage_context
|
||||
self.last_map_location = map_location
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.storage_context = None
|
||||
self.last_map_location = None
|
||||
|
||||
with set_deserialization_context():
|
||||
result = unpickler.load()
|
||||
|
||||
# TODO from zdevito:
|
||||
# This stateful weird function will need to be removed in our efforts
|
||||
# to unify the format. It has a race condition if multiple python
|
||||
# threads try to read independent files
|
||||
torch._utils._validate_loaded_sparse_tensors()
|
||||
|
||||
return result
|
||||
|
||||
def id(self):
|
||||
"""
|
||||
Returns internal identifier that torch.package uses to distinguish :class:`PackageImporter` instances.
|
||||
Looks like::
|
||||
|
||||
<torch_package_0>
|
||||
"""
|
||||
return self._mangler.parent_name()
|
||||
|
||||
def file_structure(
|
||||
self, *, include: "GlobPattern" = "**", exclude: "GlobPattern" = ()
|
||||
) -> Directory:
|
||||
"""Returns a file structure representation of package's zipfile.
|
||||
|
||||
Args:
|
||||
include (list[str] | str): An optional string e.g. ``"my_package.my_subpackage"``, or optional list of strings
|
||||
for the names of the files to be included in the zipfile representation. This can also be
|
||||
a glob-style pattern, as described in :meth:`PackageExporter.mock`
|
||||
|
||||
exclude (list[str] | str): An optional pattern that excludes files whose name match the pattern.
|
||||
|
||||
Returns:
|
||||
:class:`Directory`
|
||||
"""
|
||||
return _create_directory_from_file_list(
|
||||
self.filename, self.zip_reader.get_all_records(), include, exclude
|
||||
)
|
||||
|
||||
def python_version(self):
|
||||
"""Returns the version of python that was used to create this package.
|
||||
|
||||
Note: this function is experimental and not Forward Compatible. The plan is to move this into a lock
|
||||
file later on.
|
||||
|
||||
Returns:
|
||||
:class:`str | None` a python version e.g. 3.8.9 or None if no version was stored with this package
|
||||
"""
|
||||
python_version_path = ".data/python_version"
|
||||
return (
|
||||
self.zip_reader.get_record(python_version_path).decode("utf-8").strip()
|
||||
if self.zip_reader.has_record(python_version_path)
|
||||
else None
|
||||
)
|
||||
|
||||
def _read_extern(self):
|
||||
return (
|
||||
self.zip_reader.get_record(".data/extern_modules")
|
||||
.decode("utf-8")
|
||||
.splitlines(keepends=False)
|
||||
)
|
||||
|
||||
def _make_module(
|
||||
self, name: str, filename: str | None, is_package: bool, parent: str
|
||||
):
|
||||
mangled_filename = self._mangler.mangle(filename) if filename else None
|
||||
spec = importlib.machinery.ModuleSpec(
|
||||
name,
|
||||
self, # type: ignore[arg-type]
|
||||
origin="<package_importer>",
|
||||
is_package=is_package,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
self.modules[name] = module
|
||||
module.__name__ = self._mangler.mangle(name)
|
||||
ns = module.__dict__
|
||||
ns["__spec__"] = spec
|
||||
ns["__loader__"] = self
|
||||
ns["__file__"] = mangled_filename
|
||||
ns["__cached__"] = None
|
||||
ns["__builtins__"] = self.patched_builtins
|
||||
ns["__torch_package__"] = True
|
||||
|
||||
# Add this module to our private global registry. It should be unique due to mangling.
|
||||
if module.__name__ in _package_imported_modules:
|
||||
raise AssertionError(
|
||||
f"module {module.__name__} already exists in _package_imported_modules"
|
||||
)
|
||||
_package_imported_modules[module.__name__] = module
|
||||
|
||||
# preemptively install on the parent to prevent IMPORT_FROM from trying to
|
||||
# access sys.modules
|
||||
self._install_on_parent(parent, name, module)
|
||||
|
||||
if filename is not None:
|
||||
if mangled_filename is None:
|
||||
raise AssertionError(
|
||||
"mangled_filename must not be None when filename is set"
|
||||
)
|
||||
# preemptively install the source in `linecache` so that stack traces,
|
||||
# `inspect`, etc. work.
|
||||
if filename in linecache.cache: # type: ignore[attr-defined]
|
||||
raise AssertionError(f"filename {filename} already in linecache.cache")
|
||||
linecache.lazycache(mangled_filename, ns)
|
||||
|
||||
code = self._compile_source(filename, mangled_filename)
|
||||
exec(code, ns)
|
||||
|
||||
return module
|
||||
|
||||
def _load_module(self, name: str, parent: str):
|
||||
cur: _PathNode = self.root
|
||||
for atom in name.split("."):
|
||||
if not isinstance(cur, _PackageNode) or atom not in cur.children:
|
||||
if name in IMPLICIT_IMPORT_ALLOWLIST:
|
||||
module = self.modules[name] = importlib.import_module(name)
|
||||
return module
|
||||
raise ModuleNotFoundError(
|
||||
f'No module named "{name}" in self-contained archive "{self.filename}"'
|
||||
f" and the module is also not in the list of allowed external modules: {self.extern_modules}",
|
||||
name=name,
|
||||
)
|
||||
cur = cur.children[atom]
|
||||
if isinstance(cur, _ExternNode):
|
||||
module = self.modules[name] = importlib.import_module(name)
|
||||
|
||||
if compat_mapping := EXTERN_IMPORT_COMPAT_NAME_MAPPING.get(name):
|
||||
for old_name, new_name in compat_mapping.items():
|
||||
module.__dict__.setdefault(old_name, new_name)
|
||||
|
||||
return module
|
||||
return self._make_module(
|
||||
name,
|
||||
cur.source_file, # type: ignore[attr-defined]
|
||||
isinstance(cur, _PackageNode),
|
||||
parent,
|
||||
)
|
||||
|
||||
def _compile_source(self, fullpath: str, mangled_filename: str):
|
||||
source = self.zip_reader.get_record(fullpath)
|
||||
source = _normalize_line_endings(source)
|
||||
return compile(source, mangled_filename, "exec", dont_inherit=True)
|
||||
|
||||
# note: named `get_source` so that linecache can find the source
|
||||
# when this is the __loader__ of a module.
|
||||
def get_source(self, module_name) -> str:
|
||||
# linecache calls `get_source` with the `module.__name__` as the argument, so we must demangle it here.
|
||||
module = self.import_module(demangle(module_name))
|
||||
return self.zip_reader.get_record(demangle(module.__file__)).decode("utf-8")
|
||||
|
||||
# note: named `get_resource_reader` so that importlib.resources can find it.
|
||||
# This is otherwise considered an internal method.
|
||||
def get_resource_reader(self, fullname):
|
||||
try:
|
||||
package = self._get_package(fullname)
|
||||
except ImportError:
|
||||
return None
|
||||
if package.__loader__ is not self:
|
||||
return None
|
||||
return _PackageResourceReader(self, fullname)
|
||||
|
||||
def _install_on_parent(self, parent: str, name: str, module: types.ModuleType):
|
||||
if not parent:
|
||||
return
|
||||
# Set the module as an attribute on its parent.
|
||||
parent_module = self.modules[parent]
|
||||
if parent_module.__loader__ is self:
|
||||
setattr(parent_module, name.rpartition(".")[2], module)
|
||||
|
||||
# note: copied from cpython's import code, with call to create module replaced with _make_module
|
||||
def _do_find_and_load(self, name):
|
||||
parent = name.rpartition(".")[0]
|
||||
module_name_no_parent = name.rpartition(".")[-1]
|
||||
if parent:
|
||||
if parent not in self.modules:
|
||||
self._gcd_import(parent)
|
||||
# Crazy side-effects!
|
||||
if name in self.modules:
|
||||
return self.modules[name]
|
||||
parent_module = self.modules[parent]
|
||||
|
||||
try:
|
||||
parent_module.__path__ # type: ignore[attr-defined]
|
||||
|
||||
except AttributeError:
|
||||
# when we attempt to import a package only containing pybinded files,
|
||||
# the parent directory isn't always a package as defined by python,
|
||||
# so we search if the package is actually there or not before calling the error.
|
||||
if isinstance(
|
||||
parent_module.__loader__,
|
||||
importlib.machinery.ExtensionFileLoader,
|
||||
):
|
||||
if name not in self.extern_modules:
|
||||
msg = (
|
||||
_ERR_MSG
|
||||
+ "; {!r} is a c extension module which was not externed. C extension modules \
|
||||
need to be externed by the PackageExporter in order to be used as we do not support interning them.}."
|
||||
).format(name, name)
|
||||
raise ModuleNotFoundError(msg, name=name) from None
|
||||
if not isinstance(
|
||||
parent_module.__dict__.get(module_name_no_parent),
|
||||
types.ModuleType,
|
||||
):
|
||||
msg = (
|
||||
_ERR_MSG
|
||||
+ "; {!r} is a c extension package which does not contain {!r}."
|
||||
).format(name, parent, name)
|
||||
raise ModuleNotFoundError(msg, name=name) from None
|
||||
else:
|
||||
msg = (_ERR_MSG + "; {!r} is not a package").format(name, parent)
|
||||
raise ModuleNotFoundError(msg, name=name) from None
|
||||
|
||||
module = self._load_module(name, parent)
|
||||
|
||||
self._install_on_parent(parent, name, module)
|
||||
|
||||
return module
|
||||
|
||||
# note: copied from cpython's import code
|
||||
def _find_and_load(self, name):
|
||||
module = self.modules.get(name, _NEEDS_LOADING)
|
||||
if module is _NEEDS_LOADING:
|
||||
return self._do_find_and_load(name)
|
||||
|
||||
if module is None:
|
||||
message = f"import of {name} halted; None in sys.modules"
|
||||
raise ModuleNotFoundError(message, name=name)
|
||||
|
||||
# To handle https://github.com/pytorch/pytorch/issues/57490, where std's
|
||||
# creation of fake submodules via the hacking of sys.modules is not import
|
||||
# friendly
|
||||
if name == "os":
|
||||
self.modules["os.path"] = cast(Any, module).path
|
||||
elif name == "typing":
|
||||
if sys.version_info < (3, 13):
|
||||
self.modules["typing.io"] = cast(Any, module).io
|
||||
self.modules["typing.re"] = cast(Any, module).re
|
||||
|
||||
return module
|
||||
|
||||
def _gcd_import(self, name, package=None, level=0):
|
||||
"""Import and return the module based on its name, the package the call is
|
||||
being made from, and the level adjustment.
|
||||
|
||||
This function represents the greatest common denominator of functionality
|
||||
between import_module and __import__. This includes setting __package__ if
|
||||
the loader did not.
|
||||
|
||||
"""
|
||||
_sanity_check(name, package, level)
|
||||
if level > 0:
|
||||
name = _resolve_name(name, package, level)
|
||||
|
||||
return self._find_and_load(name)
|
||||
|
||||
# note: copied from cpython's import code
|
||||
def _handle_fromlist(self, module, fromlist, *, recursive=False):
|
||||
"""Figure out what __import__ should return.
|
||||
|
||||
The import_ parameter is a callable which takes the name of module to
|
||||
import. It is required to decouple the function from assuming importlib's
|
||||
import implementation is desired.
|
||||
|
||||
"""
|
||||
module_name = demangle(module.__name__)
|
||||
# The hell that is fromlist ...
|
||||
# If a package was imported, try to import stuff from fromlist.
|
||||
if hasattr(module, "__path__"):
|
||||
for x in fromlist:
|
||||
if not isinstance(x, str):
|
||||
if recursive:
|
||||
where = module_name + ".__all__"
|
||||
else:
|
||||
where = "``from list''"
|
||||
raise TypeError(
|
||||
f"Item in {where} must be str, not {type(x).__name__}"
|
||||
)
|
||||
elif x == "*":
|
||||
if not recursive and hasattr(module, "__all__"):
|
||||
self._handle_fromlist(module, module.__all__, recursive=True)
|
||||
elif not hasattr(module, x):
|
||||
from_name = f"{module_name}.{x}"
|
||||
try:
|
||||
self._gcd_import(from_name)
|
||||
except ModuleNotFoundError as exc:
|
||||
# Backwards-compatibility dictates we ignore failed
|
||||
# imports triggered by fromlist for modules that don't
|
||||
# exist.
|
||||
if (
|
||||
exc.name == from_name
|
||||
and self.modules.get(from_name, _NEEDS_LOADING) is not None
|
||||
):
|
||||
continue
|
||||
raise
|
||||
return module
|
||||
|
||||
def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):
|
||||
if level == 0:
|
||||
module = self._gcd_import(name)
|
||||
else:
|
||||
globals_ = globals if globals is not None else {}
|
||||
package = _calc___package__(globals_)
|
||||
module = self._gcd_import(name, package, level)
|
||||
if not fromlist:
|
||||
# Return up to the first dot in 'name'. This is complicated by the fact
|
||||
# that 'name' may be relative.
|
||||
if level == 0:
|
||||
return self._gcd_import(name.partition(".")[0])
|
||||
elif not name:
|
||||
return module
|
||||
else:
|
||||
# Figure out where to slice the module's name up to the first dot
|
||||
# in 'name'.
|
||||
cut_off = len(name) - len(name.partition(".")[0])
|
||||
# Slice end needs to be positive to alleviate need to special-case
|
||||
# when ``'.' not in name``.
|
||||
module_name = demangle(module.__name__)
|
||||
return self.modules[module_name[: len(module_name) - cut_off]]
|
||||
else:
|
||||
return self._handle_fromlist(module, fromlist)
|
||||
|
||||
def _get_package(self, package):
|
||||
"""Take a package name or module object and return the module.
|
||||
|
||||
If a name, the module is imported. If the passed or imported module
|
||||
object is not a package, raise an exception.
|
||||
"""
|
||||
if hasattr(package, "__spec__"):
|
||||
if package.__spec__.submodule_search_locations is None:
|
||||
raise TypeError(f"{package.__spec__.name!r} is not a package")
|
||||
else:
|
||||
return package
|
||||
else:
|
||||
module = self.import_module(package)
|
||||
if module.__spec__.submodule_search_locations is None:
|
||||
raise TypeError(f"{package!r} is not a package")
|
||||
else:
|
||||
return module
|
||||
|
||||
def _zipfile_path(self, package, resource=None):
|
||||
package = self._get_package(package)
|
||||
if package.__loader__ is not self:
|
||||
raise AssertionError(
|
||||
f"package.__loader__ must be self, got {package.__loader__}"
|
||||
)
|
||||
name = demangle(package.__name__)
|
||||
if resource is not None:
|
||||
resource = _normalize_path(resource)
|
||||
return f"{name.replace('.', '/')}/{resource}"
|
||||
else:
|
||||
return f"{name.replace('.', '/')}"
|
||||
|
||||
def _get_or_create_package(self, atoms: list[str]) -> "_PackageNode | _ExternNode":
|
||||
cur = self.root
|
||||
for i, atom in enumerate(atoms):
|
||||
node = cur.children.get(atom, None)
|
||||
if node is None:
|
||||
node = cur.children[atom] = _PackageNode(None)
|
||||
if isinstance(node, _ExternNode):
|
||||
return node
|
||||
if isinstance(node, _ModuleNode):
|
||||
name = ".".join(atoms[:i])
|
||||
raise ImportError(
|
||||
f"inconsistent module structure. module {name} is not a package, but has submodules"
|
||||
)
|
||||
if not isinstance(node, _PackageNode):
|
||||
raise AssertionError(
|
||||
f"expected _PackageNode, got {type(node).__name__}"
|
||||
)
|
||||
cur = node
|
||||
return cur
|
||||
|
||||
def _add_file(self, filename: str):
|
||||
"""Assembles a Python module out of the given file. Will ignore files in the .data directory.
|
||||
|
||||
Args:
|
||||
filename (str): the name of the file inside of the package archive to be added
|
||||
"""
|
||||
*prefix, last = filename.split("/")
|
||||
if len(prefix) > 1 and prefix[0] == ".data":
|
||||
return
|
||||
package = self._get_or_create_package(prefix)
|
||||
if isinstance(package, _ExternNode):
|
||||
raise ImportError(
|
||||
f"inconsistent module structure. package contains a module file {filename}"
|
||||
f" that is a subpackage of a module marked external."
|
||||
)
|
||||
if last == "__init__.py":
|
||||
package.source_file = filename
|
||||
elif last.endswith(".py"):
|
||||
package_name = last[: -len(".py")]
|
||||
package.children[package_name] = _ModuleNode(filename)
|
||||
|
||||
def _add_extern(self, extern_name: str):
|
||||
*prefix, last = extern_name.split(".")
|
||||
package = self._get_or_create_package(prefix)
|
||||
if isinstance(package, _ExternNode):
|
||||
return # the shorter extern covers this extern case
|
||||
package.children[last] = _ExternNode()
|
||||
|
||||
|
||||
_NEEDS_LOADING = object()
|
||||
_ERR_MSG_PREFIX = "No module named "
|
||||
_ERR_MSG = _ERR_MSG_PREFIX + "{!r}"
|
||||
|
||||
|
||||
class _PathNode:
|
||||
__slots__ = []
|
||||
|
||||
|
||||
class _PackageNode(_PathNode):
|
||||
def __init__(self, source_file: str | None):
|
||||
self.source_file = source_file
|
||||
self.children: dict[str, _PathNode] = {}
|
||||
|
||||
|
||||
class _ModuleNode(_PathNode):
|
||||
__slots__ = ["source_file"]
|
||||
|
||||
def __init__(self, source_file: str):
|
||||
self.source_file = source_file
|
||||
|
||||
|
||||
class _ExternNode(_PathNode):
|
||||
pass
|
||||
|
||||
|
||||
# A private global registry of all modules that have been package-imported.
|
||||
_package_imported_modules: WeakValueDictionary = WeakValueDictionary()
|
||||
|
||||
# `inspect` by default only looks in `sys.modules` to find source files for classes.
|
||||
# Patch it to check our private registry of package-imported modules as well.
|
||||
_orig_getfile = inspect.getfile
|
||||
|
||||
|
||||
def _patched_getfile(object):
|
||||
if inspect.isclass(object):
|
||||
if object.__module__ in _package_imported_modules:
|
||||
return _package_imported_modules[object.__module__].__file__
|
||||
return _orig_getfile(object)
|
||||
|
||||
|
||||
inspect.getfile = _patched_getfile
|
||||
|
||||
|
||||
class _PackageResourceReader:
|
||||
"""Private class used to support PackageImporter.get_resource_reader().
|
||||
|
||||
Confirms to the importlib.abc.ResourceReader interface. Allowed to access
|
||||
the innards of PackageImporter.
|
||||
"""
|
||||
|
||||
def __init__(self, importer, fullname):
|
||||
self.importer = importer
|
||||
self.fullname = fullname
|
||||
|
||||
def open_resource(self, resource):
|
||||
from io import BytesIO
|
||||
|
||||
return BytesIO(self.importer.load_binary(self.fullname, resource))
|
||||
|
||||
def resource_path(self, resource):
|
||||
# The contract for resource_path is that it either returns a concrete
|
||||
# file system path or raises FileNotFoundError.
|
||||
if isinstance(
|
||||
self.importer.zip_reader, DirectoryReader
|
||||
) and self.importer.zip_reader.has_record(
|
||||
os.path.join(self.fullname, resource)
|
||||
):
|
||||
return os.path.join(
|
||||
self.importer.zip_reader.directory, self.fullname, resource
|
||||
)
|
||||
raise FileNotFoundError
|
||||
|
||||
def is_resource(self, name):
|
||||
path = self.importer._zipfile_path(self.fullname, name)
|
||||
return self.importer.zip_reader.has_record(path)
|
||||
|
||||
def contents(self):
|
||||
from pathlib import Path
|
||||
|
||||
filename = self.fullname.replace(".", "/")
|
||||
|
||||
fullname_path = Path(self.importer._zipfile_path(self.fullname))
|
||||
files = self.importer.zip_reader.get_all_records()
|
||||
subdirs_seen = set()
|
||||
for filename in files:
|
||||
try:
|
||||
relative = Path(filename).relative_to(fullname_path)
|
||||
except ValueError:
|
||||
continue
|
||||
# If the path of the file (which is relative to the top of the zip
|
||||
# namespace), relative to the package given when the resource
|
||||
# reader was created, has a parent, then it's a name in a
|
||||
# subdirectory and thus we skip it.
|
||||
parent_name = relative.parent.name
|
||||
if len(parent_name) == 0:
|
||||
yield relative.name
|
||||
elif parent_name not in subdirs_seen:
|
||||
subdirs_seen.add(parent_name)
|
||||
yield parent_name
|
||||
Reference in New Issue
Block a user