import os import hashlib import re from typing import Tuple, List from .chunk_store import ChunkStore from .chunk import ChunkID, PositionFileRange, PositionType from ..object import HashValue from .tracker import ChunkTracker from .chunk_list import ChunkList class ChunkListWriter: def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker): self.chunk_store = chunk_store self.chunk_tracker = chunk_tracker def create_chunk_list_from_file( self, file_path: str, chunk_size: int, restore: bool ) -> ChunkList: assert ( chunk_size % (1024 * 1024) == 0 ), "chunk size should be an integral multiple of 1MB" chunk_list = [] hash_obj = hashlib.sha256() with open(file_path, "rb") as file: while True: chunk = file.read(chunk_size) if not chunk: break chunk_len = len(chunk) chunk_id = ChunkID.hash_data(chunk) chunk_list.append(chunk_id) hash_obj.update(chunk) if restore: self.chunk_tracker.add_position( chunk_id, file_path, PositionType.ChunkStore ) self.chunk_store.put_chunk(chunk_id, chunk) else: pos = file.tell() file_range = PositionFileRange( file_path, pos - chunk_len, pos ) self.chunk_tracker.add_position( chunk_id, str(file_range), PositionType.FileRange ) file_hash = HashValue(hash_obj.digest()) # print(f"calc file hash: {file_path}, {file_hash}") return ChunkList(chunk_list, file_hash) def create_chunk_list_from_text( self, text: str, chunk_max_words: int, separator_chars: str = ".," ) -> ChunkList: text_list = self._split_text_list(text, chunk_max_words, separator_chars) chunk_list = [] hash_obj = hashlib.sha256() for text in text_list: chunk_bytes = text.encode("utf-8") hash_obj.update(chunk_bytes) chunk_id = ChunkID.hash_data(chunk_bytes) chunk_list.append(chunk_id) self.chunk_tracker.add_position(chunk_id, "", PositionType.ChunkStore) self.chunk_store.put_chunk(chunk_id, chunk_bytes) hash = HashValue(hash_obj.digest()) return ChunkList(chunk_list, hash) @staticmethod def _split_text_list( text: str, chunk_max_words: int, separator_chars: str = ".," ) -> List[str]: sentences = re.split(f"[{separator_chars}]", text) # chunk_list = [] # chunk = [] # word_count = 0 # for sentence in sentences: # words = sentence.split() # for word in words: # if word_count < chunk_max_words: # chunk.append(word) # word_count += 1 # else: # chunk_list.append(" ".join(chunk)) # chunk = [word] # word_count = 1 # if chunk: # chunk_list.append(" ".join(chunk)) # return chunk_list return sentences