fix: doctests

i2mint · Jan 3, 2024 · ffadb09 · ffadb09
1 parent acfd50d
commit ffadb09
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 24 deletions.
diff --git a/README.md b/README.md
@@ -53,20 +53,42 @@ would be like this:
     ...     'documents': ['contents for piece', 'contents for of'],
     ...     'metadatas': [{'author': 'me'}, {'author': 'you'}],
     ... }
-    >>> list(collection)
-    ['piece', 'of']
-    >>>
-    >>> assert collection[['piece', 'of']] == {
-    ...     'ids': ['piece', 'of'],
+
+Now we have two documents in the collection:
+
+    >>> len(collection)
+    2
+
+Note, though, that the order of the documents is not guaranteed.
+
+    >>> sorted(collection)
+    ['of', 'piece']
+
+    >>> assert collection['piece'] == {
+    ...     'ids': ['piece'],
     ...     'embeddings': None,
-    ...     'metadatas': [{'author': 'me'}, {'author': 'you'}],
-    ...     'documents': ['contents for piece', 'contents for of'],
+    ...     'metadatas': [{'author': 'me'}],
+    ...     'documents': ['contents for piece'],
     ...     'uris': None,
-    ...     'data': None,
+    ...     'data': None
     ... }
 
+    >>> assert collection['of'] == {
+    ...     'ids': ['of'],
+    ...     'embeddings': None,
+    ...     'metadatas': [{'author': 'you'}],
+    ...     'documents': ['contents for of'],
+    ...     'uris': None,
+    ...     'data': None
+    ... }
+
+You can also read multiple documents at once.
+But note that the order of the documents is not guaranteed.
+
+    >>> collection[['piece', 'of']] == collection[['of', 'piece']]
+    True
 
-But you can read or write one document at a time too.
+You can read or write one document at a time too.
 
     >>> collection['cake'] = {
     ...     "documents": "contents for cake",

diff --git a/chromadol/__init__.py b/chromadol/__init__.py
@@ -48,20 +48,42 @@
 ...     'documents': ['contents for piece', 'contents for of'],
 ...     'metadatas': [{'author': 'me'}, {'author': 'you'}],
 ... }
+
+Now we have two documents in the collection:
+
+>>> len(collection)
+2
+
+Note, though, that the order of the documents is not guaranteed.
+
 >>> sorted(collection)
 ['of', 'piece']
->>>
->>> assert collection[['piece', 'of']] == {
-...     'ids': ['piece', 'of'],
+
+>>> assert collection['piece'] == {
+...     'ids': ['piece'],
 ...     'embeddings': None,
-...     'metadatas': [{'author': 'me'}, {'author': 'you'}],
-...     'documents': ['contents for piece', 'contents for of'],
+...     'metadatas': [{'author': 'me'}],
+...     'documents': ['contents for piece'],
 ...     'uris': None,
-...     'data': None,
+...     'data': None
 ... }
 
+>>> assert collection['of'] == {
+...     'ids': ['of'],
+...     'embeddings': None,
+...     'metadatas': [{'author': 'you'}],
+...     'documents': ['contents for of'],
+...     'uris': None,
+...     'data': None
+... }
+
+You can also read multiple documents at once.
+But note that the order of the documents is not guaranteed.
+
+>>> collection[['piece', 'of']] == collection[['of', 'piece']]
+True
 
-But you can read or write one document at a time too.
+You can read or write one document at a time too.
 
 >>> collection['cake'] = {
 ...     "documents": "contents for cake",

diff --git a/chromadol/tests/base_test.py b/chromadol/tests/base_test.py
@@ -22,23 +22,40 @@ def test_simple():
     # note that just accessing the collection creates it (by default)
     assert list(client) == ['chromadol_test']
     assert list(collection) == []
+    assert len(collection) == 0
 
     # chromadb is designed to operate on multiple documents at once, so
     # specifying it's keys and contents (and any extras) list this:
     collection[['piece', 'of']] = {
         'documents': ['contents for piece', 'contents for of'],
         'metadatas': [{'author': 'me'}, {'author': 'you'}],
     }
+
+    # Now we have two documents in the collection:
+
+    assert len(collection) == 2
+
+    # Note, though, that the order of the documents is not guaranteed.
+
     assert sorted(collection) == ['of', 'piece']
 
-    assert collection[['piece', 'of']] == {
-        'ids': ['piece', 'of'],
-        'embeddings': None,
-        'metadatas': [{'author': 'me'}, {'author': 'you'}],
-        'documents': ['contents for piece', 'contents for of'],
-        'uris': None,
-        'data': None,
-    }
+    assert collection['piece'] == {'ids': ['piece'],
+    'embeddings': None,
+    'metadatas': [{'author': 'me'}],
+    'documents': ['contents for piece'],
+    'uris': None,
+    'data': None}
+
+    assert collection['of'] == {'ids': ['of'],
+    'embeddings': None,
+    'metadatas': [{'author': 'you'}],
+    'documents': ['contents for of'],
+    'uris': None,
+    'data': None}
+
+    # You can also read multiple documents at once.
+    # But note that the order of the documents is not guaranteed.
+    assert collection[['piece', 'of']] == collection[['of', 'piece']]
 
     # But you can read or write one document at a time too.
     collection['cake'] = {