| import os.path |
| from glob import glob |
| from typing import cast |
|
|
| import torch |
| from torch.types import Storage |
|
|
| |
| class _HasStorage(object): |
| def __init__(self, storage): |
| self._storage = storage |
|
|
| def storage(self): |
| return self._storage |
|
|
|
|
| class DirectoryReader(object): |
| """ |
| 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 = [] |
| for filename in glob(f"{self.directory}/**", recursive=True): |
| if not os.path.isdir(filename): |
| files.append(filename[len(self.directory) + 1 :]) |
| return files |
|
|
