| 24 | |
| 25 | |
| 26 | class DocumentDatabase: |
| 27 | def __init__(self, temp_dir): |
| 28 | self.temp_dir = TemporaryDirectory(dir=temp_dir) |
| 29 | self.working_dir = Path(self.temp_dir.name) |
| 30 | self.document_shelf_filepath = self.working_dir / 'shelf.db' |
| 31 | self.document_shelf = shelve.open(str(self.document_shelf_filepath), |
| 32 | flag='n', protocol=-1) |
| 33 | self.docids = [] |
| 34 | |
| 35 | def add_document(self, doc_id, document): |
| 36 | self.document_shelf[str(doc_id)] = document |
| 37 | self.docids.append(doc_id) |
| 38 | |
| 39 | def __len__(self): |
| 40 | return len(self.docids) |
| 41 | |
| 42 | def __getitem__(self, item): |
| 43 | return self.document_shelf[str(item)] |
| 44 | |
| 45 | def __enter__(self): |
| 46 | return self |
| 47 | |
| 48 | def __exit__(self, exc_type, exc_val, traceback): |
| 49 | if self.document_shelf is not None: |
| 50 | self.document_shelf.close() |
| 51 | if self.temp_dir is not None: |
| 52 | self.temp_dir.cleanup() |
| 53 | |
| 54 | |
| 55 | |