fromfsspec.asynimportAsyncFileSystemimportasyncioimportposixpathfromtypingimportList,OptionalimportstructclassReadOnlyZipFileSystem(AsyncFileSystem):"""An async read-only file system for uncompressed zip files using fsspec.    Mounts an uncompressed zip file as a read-only file system. Supports ZIP64.    Only supports ZIP_STORED (uncompressed). Multi-disk zip files are not    supported. Written for reading a zipped Zarr from S3 storage and therefore    only includes methods used by Zarr.    Reads 64KB from the end of the zip file to capture the ZIP64 end of central    directory (EOCD), the ZIP64 EOCD locator, and the standard EOCD, using    _cat_file with negative start offset and therefore not needing to query the    file size first. Reads the CD using _cat_file with positive start offset.    Parses the CD for the file names and the header offsets. Requires that the    CD contains an entry of each directory (except for the autogenerated root    dir) before the entries of their files and subdirectories. Requires that    the CD entries appear in the same order as the file headers and the files.    Assumes that there are no gaps between file headers and files. This way the    file headers need not be read, because the file offset will be the offset    of the next file header (or CD for the last file) minus the file size. File    datetimes are not available.    While not currently supported, support for compressed zip files could be    implemented by reading the file headers (ending the read at the next header    or CD) and by decompression.    """protocol="zipfs"MAX_ZIP_TAIL_READ=64*1024def__init__(self,fs:AsyncFileSystem,path:str,**kwargs):"""Initialize the ReadOnlyZipFileSystem.        Args:            fs: The underlying AsyncFileSystem containing the zip file.            path: Path to the zip file in the underlying file system.            **kwargs: Additional arguments passed to AsyncFileSystem.        """super().__init__(**kwargs)self.asynchronous=Trueself.fs=fsself.path=pathself._files=Noneself._lock=asyncio.Lock()asyncdef_initialize(self):"""Initialize self._files by reading and parsing the central directory.        All other methods that require self._files first await this method.        They never modify self._files. Locking is used to ensure that the        initialization is thread-safe. The other methods need not lock        explicitly.        """asyncwithself._lock:ifself._filesisnotNone:return# Read tail of file (up to MAX_ZIP_TAIL_READ) from the enddata=awaitself.fs._cat_file(self.path,start=-self.MAX_ZIP_TAIL_READ,end=None)iflen(data)<22:# Minimum size for standard EOCDraiseValueError(f"EOCD doesn't fit in{self.path}:{len(data)} bytes")# EOCD variables and their lengths (order matters)# These dicts use exact wording from https://pkwaredownloads.blob.core.windows.net/pkware-general/Documentation/APPNOTE-6.3.9.TXTeocd_var_lengths= {'end of central dir signature':4,# Signature (0x06054b50)'number of this disk':2,# Disk number'number of the disk with the start of the central directory':2,# Disk number of the central directory'total number of entries in the central directory on this disk':2,# Number of entries on this disk'total number of entries in the central directory':2,# Total number of entries'size of the central directory offset of start of central directory with respect to':4,# Size of the central directory'the starting disk number':4,# Offset of the start of the central directory'.ZIP file comment length':2# Length of the comment            }# ZIP64 EOCD variables and their lengths (order matters)zip64_eocd_var_lengths= {'zip64 end of central dir signature':4,'size of zip64 end of central directory record':8,'version made by':2,'version needed to extract':2,'number of this disk':4,'number of the disk with the start of the central directory':4,'total number of entries in the central directory on this disk':8,'total number of entries in the central directory':8,'size of the central directory':8,'offset of start of central directory with respect to the starting disk number':8            }# Map lengths to struct formatsvar_length_to_format= {2:'H',# Unsigned short (2 bytes)4:'L',# Unsigned long (4 bytes)8:'Q'# Unsigned long long (8 bytes)            }# Parse EOCDeocd= {}is_zip64=Falseeocd_pos=data[:-20+4].rfind(b'\x50\x4b\x05\x06')# 20 bytes for EOCD, 4 bytes for EOCD signatureifeocd_pos==-1:raiseValueError(f"No EOCD in the last{self.MAX_ZIP_TAIL_READ} bytes of{self.path}")pos=eocd_posforvar,lengthineocd_var_lengths.items():eocd[var]=struct.unpack_from(f'<{var_length_to_format[length]}',data,pos)[0]# Check for ZIP64 values 0xFF or 0xFFFFFFFFifeocd[var]==2**(length*8)-1:eocd[var]=Noneis_zip64=Truepos+=lengthifis_zip64:iflen(data)-22<56+20:# 56 bytes for ZIP64 EOCD + 20 bytes for locatorraiseValueError(f"ZIP64 EOCD and ZIP64 EOCD locator do not fit in{self.path}")# Find ZIP64 EOCDzip64_eocd_pos=data[:eocd_pos-56+4].rfind(b'\x50\x4b\x06\x06')# 20 bytes for EOCD, 56 bytes for ZIP64 EOCD, 4 bytes for ZIP64 EOCD signatureifzip64_eocd_pos==-1:raiseValueError(f"No ZIP64 EOCD in the last{self.MAX_ZIP_TAIL_READ} bytes of{self.path}")pos=zip64_eocd_posforvar,lengthinzip64_eocd_var_lengths.items():eocd[var]=struct.unpack_from(f'<{var_length_to_format[length]}',data,pos)[0]pos+=length# Require single-disk zipifeocd['number of this disk']!=0oreocd['number of the disk with the start of the central directory']!=0oreocd['total number of entries in the central directory on this disk']!=eocd['total number of entries in the central directory']:raiseValueError(f"Unsupported multi-disk central directory in{self.path}")# Convenience variablescd_size=eocd['size of the central directory']cd_offset=eocd['offset of start of central directory with respect to the starting disk number']cd_entries=eocd['total number of entries in the central directory']# Read and parse central directoryifcd_size==0:# No central directory, empty zip filereturncd_data=awaitself.fs._cat_file(self.path,start=cd_offset,end=cd_offset+cd_size)# Save data to fileiflen(cd_data)!=cd_size:raiseValueError(f"Failed to read central directory: expected{cd_size} bytes, got{len(cd_data)}")# Central directory file header variables and their lengthscd_file_header_var_lengths= {'central file header signature':4,'version made by':2,'version needed to extract':2,'general purpose bit flag':2,'compression method':2,'last mod file time':2,'last mod file date':2,'crc-32':4,'compressed size':4,'uncompressed size':4,'file name length':2,'extra field length':2,'file comment length':2,'disk number start':2,'internal file attributes':2,'external file attributes':4,'relative offset of local header':4            }# Central or local file header ZIP64 extended information extra field variables and their lengths# These share names with the standard fields but have different lengthszip64_file_header_var_lengths= {'uncompressed size':8,'compressed size':8,'relative offset of local header':8,'disk number start':4            }# Autocreate root dirself._files= {'': {'children': []                }            }# Parse central directory entriespos=0previous_file=Noneforfile_indexinrange(cd_entries):ifpos+46>len(cd_data):# 46 bytes for the central directory file headerraiseValueError(f"Truncated central directory entry in{self.path}")cd_file_header= {}forvar,lengthincd_file_header_var_lengths.items():cd_file_header[var]=struct.unpack_from(f'<{var_length_to_format[length]}',cd_data,pos)[0]# Check for ZIP64 values 0xFF or 0xFFFFFFFFifcd_file_header[var]==2**(length*8)-1:cd_file_header[var]=Nonepos+=lengthifcd_file_header['central file header signature']!=0x02014b50:raiseValueError(f"Invalid central directory header signature in{self.path}")ifcd_file_header['compression method']!=0orcd_file_header['compressed size']!=cd_file_header['uncompressed size']:raiseValueError(f"File in{self.path} is not stored (uncompressed)")utf8=cd_file_header['general purpose bit flag']&0x800!=0# Bit 11# Convenience variablesfname_len=cd_file_header['file name length']extra_len=cd_file_header['extra field length']comment_len=cd_file_header['file comment length']# Read filenameifpos+fname_len>len(cd_data):raiseValueError(f"Truncated filename in{self.path}")ifutf8:fname=cd_data[pos:pos+fname_len].decode('utf-8')else:fname=cd_data[pos:pos+fname_len].decode('ascii')pos+=fname_len# Parse extra fieldextra_end=pos+extra_lenifextra_end>len(cd_data):raiseValueError(f"Truncated extra field in{self.path}")whilepos<extra_end:ifpos+4>extra_end:raiseValueError(f"Truncated extra field in{self.path}")tag,size=struct.unpack_from('<HH',cd_data,pos)pos+=4ifpos+size>extra_end:raiseValueError(f"Truncated extra field in{self.path}")backup_pos=posiftag==0x0001:# ZIP64 extended information extra fieldforvar,lengthinzip64_file_header_var_lengths.items():# Only parse variables that were marked as ZIP64ifcd_file_header[var]isNone:cd_file_header[var]=struct.unpack_from(f'<{var_length_to_format[length]}',cd_data,pos)[0]pos+=lengthelse:iftag==0x5455:# Extended Timestamp, not handledNoneeliftag==0x7875:# Info-ZIP New Unix Extra Field, skip itNoneelse:# Unknown extra field, skip itNonepos+=sizeifbackup_pos+size!=pos:raiseValueError(f"Invalid extra field size in{self.path}")pos=backup_pos+size# Convenience variablessize=cd_file_header['compressed size']offset=cd_file_header['relative offset of local header']# Skip commentpos+=comment_len# Detect directoryis_dir=fname.endswith('/')ifis_dir:# Remove trailing slashfname=fname[:-1]# Store file infoifis_dir:# Dirs can be recognized by key 'children' existingfile= {'children': []                    }else:# It's a filefile= {'size':size,'offset':offset                    }self._files[fname]=file# Add file as the parent directory's childiffname!='':parent=posixpath.dirname(fname)ifparentnotinself._files:raiseNotImplementedError("Autocreation of parent folder {parent} not implemented, in {self.path}")self._files[parent]['children'].append(fname)# Calculate offset of previous fileifprevious_fileisnotNoneand'children'notinprevious_file:# We require the files to be stored in the order of the central directory entries,# although sorting would be an optionifoffset<=previous_file['offset']:raiseValueError(f"Non-ascending order of local header offsets in{self.path}")# Previous file's offset = current file's offset - previous file's size# This should work unless there is empty space in the zip file, which is unlikely.previous_file['offset']=offset-previous_file['size']previous_file=file# Calculate offset of previous fileifprevious_fileisnotNoneand'children'notinprevious_file:ifcd_offset<=previous_file['offset']:raiseValueError(f"Non-ascending order of local header offsets in{self.path}")# Last file ends where central directory beginsprevious_file['offset']=cd_offset-previous_file['size']asyncdef_ls(self,path:str,detail:bool=True,**kwargs)->List:""" List files and directories in the given path.        If the path points to a file, list just the file.        """# Always await self._initialize() in functions needing self._filesawaitself._initialize()# Internally we don't use a root slash, so strip it. Also strip any trailing slash.path=posixpath.normpath(path).lstrip('/').rstrip('/')# Helper function to get file name or detailsdefget_file_listing(file,fname,detail):ifdetail:if'children'infile:return {'name':f'/{fname}','type':'directory','size':0,'created':None,'islink':False                                    }else:return {'name':f'/{fname}','type':'file','size':file['size'],'created':None,'islink':False                    }else:returnf'/{fname}'# List childrenresults= []ifpathnotinself._files:raiseFileNotFoundError(f"Path{path} not found")file=self._files[path]if'children'infile:# Path points to a dirforchildinfile['children']:results.append(get_file_listing(self._files[child],child,detail))else:# Path points to a fileresults= [get_file_listing(file,path,detail)]returnresultsasyncdef_cat_file(self,path:str,start:Optional[int]=None,end:Optional[int]=None,**kwargs)->bytes:"""Read the contents of a file in the zip."""# Always await self._initialize() in functions needing self._filesawaitself._initialize()# Internally we don't use a root slash, so strip it. Also strip any trailing slash.path=posixpath.normpath(path).lstrip('/').rstrip('/')# Check if the file is availableifpathnotinself._files:raiseFileNotFoundError(f"File{path} not found")elif'children'inself._files[path]:raiseFileNotFoundError(f"{path} is a directory")# Get offset and size of the file in the zip fileinfo=self._files[path]offset=info['offset']size=info['size']# Set start to beginning of file if not specifiedstart=startor0# Convert negative start (relative to end of file) to positive startifstart<0:start=max(0,size+start)# Clamp too large negative start to the beginning of file# Set end to end of file if not specifiedend=endorsize# Convert negative start (relative to end of file) to positive startifend<0:end=max(0,size+end)# Clamp too large negative start to the beginning of file# For start beyond the end of the file or the end, return empty dataifstart>=sizeorend<=start:returnb''# Calculate zip file read start and read sizeread_start=offset+startread_size=min(end,size)-start# Clamp too large end at size# Read datadata=awaitself.fs._cat_file(self.path,start=read_start,end=read_start+read_size)returndata