Skip to content

Commit

Permalink
Allow Table.get by doc_id to have O(1) behavior
Browse files Browse the repository at this point in the history
  • Loading branch information
richardsheridan committed Feb 15, 2022
1 parent e2efe98 commit d082bc6
Showing 1 changed file with 8 additions and 27 deletions.
35 changes: 8 additions & 27 deletions tinydb/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -289,7 +289,7 @@ def get(
if doc_id is not None:
# Retrieve a document specified by its ID
table = self._read_table()
raw_doc = table.get(doc_id, None)
raw_doc = table.get(str(doc_id), None)

if raw_doc is None:
return None
Expand Down Expand Up @@ -610,20 +610,7 @@ def __len__(self):
Count the total number of documents in this table.
"""

# Using self._read_table() will convert all documents into
# the document class. But for counting the number of documents
# this conversion is not necessary, thus we read the storage
# directly here

tables = self._storage.read()

if tables is None:
return 0

try:
return len(tables[self.name])
except KeyError:
return 0
return len(self._read_table())

def __iter__(self) -> Iterator[Document]:
"""
Expand All @@ -635,7 +622,7 @@ def __iter__(self) -> Iterator[Document]:
# Iterate all documents and their IDs
for doc_id, doc in self._read_table().items():
# Convert documents to the document class
yield self.document_class(doc, doc_id)
yield self.document_class(doc, self.document_id_class(doc_id))

def _get_next_id(self):
"""
Expand Down Expand Up @@ -672,14 +659,13 @@ def _get_next_id(self):

return next_id

def _read_table(self) -> Dict[int, Mapping]:
def _read_table(self) -> Dict[str, Mapping]:
"""
Read the table data from the underlying storage.
Here we read the data from the underlying storage and convert all
IDs to the document ID class. Documents themselves are NOT yet
transformed into the document class, we may not want to convert
*all* documents when returning only one document for example.
Documents and doc_ids are NOT yet transformed, as
we may not want to convert *all* documents when returning
only one document for example.
"""

# Retrieve the tables from the storage
Expand All @@ -696,12 +682,7 @@ def _read_table(self) -> Dict[int, Mapping]:
# The table does not exist yet, so it is empty
return {}

# Convert all document IDs to the correct document ID class and return
# the table data dict
return {
self.document_id_class(doc_id): doc
for doc_id, doc in table.items()
}
return table

def _update_table(self, updater: Callable[[Dict[int, Mapping]], None]):
"""
Expand Down

0 comments on commit d082bc6

Please sign in to comment.