U <ca@sXddlZddlmZddlmZddlZddlmZGdddeZ GdddeZ dS) N)glob)cast)Storagec@seZdZddZddZdS) _HasStoragecCs ||_dSNZ_storage)selfstorager C/tmp/pip-unpacked-wheel-gikjz4vx/torch/package/_directory_reader.py__init__ sz_HasStorage.__init__cCs|jSrr)rr r r r sz_HasStorage.storageN)__name__ __module__ __qualname__r r r r r r r src@s8eZdZdZddZddZddZdd Zd d Zd S) DirectoryReaderar 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. cCs ||_dSr) directory)rrr r r r szDirectoryReader.__init__c Cs:|jd|}t|d}|W5QRSQRXdS)N/rb)ropenread)rnamefilenamefr r r get_records zDirectoryReader.get_recordcCs>|jd|}tj||}tttj}t|j||dS)Nr)rnbytes) rtorch_utilsZ _element_sizerrZUntypedStorager from_file)rrZnumelZdtyperrr r r r get_storage_from_record#s z'DirectoryReader.get_storage_from_recordcCstj|j|}tj|Sr)ospathjoinrisfile)rr full_pathr r r has_record)szDirectoryReader.has_recordcCsJg}t|jdddD],}tj|s||t|jddq|S)Nz/**T) recursive)rrrr isdirappendlen)rfilesrr r r get_all_records-s  zDirectoryReader.get_all_recordsN) r rr__doc__r rrr$r+r r r r rs  r) os.pathrrtypingrrZ torch.typesrobjectrrr r r r s