Implemented container methods (#86)

* implemented container methods

* changed relative to absolute imports

* fix flake8

* revert absolute import change

testbook/client.py

@@ -289,9 +289,6 @@ def value(self, code: str) -> Any:
             e.save_varname = save_varname
             raise e
 
-    def _eq_in_notebook(self, lhs: str, rhs: Any) -> bool:
-        return self.value("{lhs} == {rhs}".format(lhs=lhs, rhs=PythonTranslator.translate(rhs)))
-
     @contextmanager
     def patch(self, target, **kwargs):
         """Used as contextmanager to patch objects in the kernel"""

testbook/reference.py

@@ -2,7 +2,10 @@
     TestbookExecuteResultNotFoundError,
     TestbookAttributeError,
     TestbookSerializeError,
+    TestbookRuntimeError
 )
+from .utils import random_varname
+from .translators import PythonTranslator
 
 
 class TestbookObjectReference:
@@ -24,7 +27,57 @@ def __getattr__(self, name):
         raise TestbookAttributeError(f"'{self._type}' object has no attribute {name}")
 
     def __eq__(self, rhs):
-        return self.tb._eq_in_notebook(self.name, rhs)
+        return self.tb.value(
+            "{lhs} == {rhs}".format(lhs=self.name, rhs=PythonTranslator.translate(rhs))
+        )
+
+    def __len__(self):
+        return self.tb.value(f"len({self.name})")
+
+    def __iter__(self):
+        iterobjectname = f"___iter_object_{random_varname()}"
+        self.tb.inject(f"""
+            {iterobjectname} = iter({self.name})
+        """)
+        return TestbookObjectReference(self.tb, iterobjectname)
+
+    def __next__(self):
+        try:
+            return self.tb.value(f"next({self.name})")
+        except TestbookRuntimeError as e:
+            if e.eclass is StopIteration:
+                raise StopIteration
+            else:
+                raise
+
+    def __getitem__(self, key):
+        try:
+            return self.tb.value(f"{self.name}.__getitem__({PythonTranslator.translate(key)})")
+        except TestbookRuntimeError as e:
+            if e.eclass is TypeError:
+                raise TypeError(e.evalue)
+            elif e.eclass is IndexError:
+                raise IndexError(e.evalue)
+            else:
+                raise
+
+    def __setitem__(self, key, value):
+        try:
+            return self.tb.inject("{name}[{key}] = {value}".format(
+                name=self.name,
+                key=PythonTranslator.translate(key),
+                value=PythonTranslator.translate(value)
+            ), pop=True)
+        except TestbookRuntimeError as e:
+            if e.eclass is TypeError:
+                raise TypeError(e.evalue)
+            elif e.eclass is IndexError:
+                raise IndexError(e.evalue)
+            else:
+                raise
+
+    def __contains__(self, item):
+        return self.tb.value(f"{self.name}.__contains__({PythonTranslator.translate(item)})")
 
     def __call__(self, *args, **kwargs):
         code = self.tb._construct_call_code(self.name, args, kwargs)

testbook/tests/resources/datamodel.ipynb

@@ -0,0 +1,43 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "political-plaintiff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mylist = [1, 2, 3, 4, 5]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "postal-contemporary",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

testbook/tests/test_datamodel.py

@@ -0,0 +1,83 @@
+import pytest
+
+from ..testbook import testbook
+
+
+@pytest.fixture(scope='module')
+def notebook():
+    with testbook('testbook/tests/resources/datamodel.ipynb', execute=True) as tb:
+        yield tb
+
+
+def test_len(notebook):
+    mylist = notebook.ref("mylist")
+
+    assert len(mylist) == 5
+
+
+def test_iter(notebook):
+    mylist = notebook.ref("mylist")
+
+    expected = []
+    for x in mylist:
+        expected.append(x)
+
+    assert mylist == expected
+
+
+def test_getitem(notebook):
+    mylist = notebook.ref("mylist")
+    mylist.append(6)
+
+    assert mylist[-1] == 6
+    assert mylist.__getitem__(-1) == 6
+
+
+def test_getitem_raisesIndexError(notebook):
+    mylist = notebook.ref("mylist")
+
+    with pytest.raises(IndexError):
+        mylist[100]
+
+
+def test_getitem_raisesTypeError(notebook):
+    mylist = notebook.ref("mylist")
+
+    with pytest.raises(TypeError):
+        mylist['hello']
+
+
+def test_setitem(notebook):
+    notebook.inject("mydict = {'key1': 'value1', 'key2': 'value1'}")
+    mydict = notebook.ref("mydict")
+
+    mydict['key3'] = 'value3'
+    assert mydict['key3'] == 'value3'
+
+    mylist = notebook.ref("mylist")
+    mylist[2] = 10
+    assert mylist[2] == 10
+
+
+def test_setitem_raisesIndexError(notebook):
+    mylist = notebook.ref("mylist")
+
+    with pytest.raises(IndexError):
+        mylist.__setitem__(10, 100)
+
+
+def test_setitem_raisesTypeError(notebook):
+    mylist = notebook.ref("mylist")
+
+    with pytest.raises(TypeError):
+        mylist.__setitem__('key', 10)
+
+
+def test_contains(notebook):
+    notebook.inject("mydict = {'key1': 'value1', 'key2': 'value1'}")
+    mydict = notebook.ref("mydict")
+
+    assert 'key1' in mydict
+    assert 'key2' in mydict
+    assert mydict.__contains__('key1')
+    assert mydict.__contains__('key2')

testbook/tests/test_reference.py

@@ -51,7 +51,7 @@ def test_reference(notebook):
     assert repr(f) == "\"<Foo value='bar'>\""
 
     # Valid attribute access
-    assert f.say_hello
+    assert f.say_hello()
 
     # Invalid attribute access
     with pytest.raises(TestbookAttributeError):

testbook/translators.py

@@ -2,8 +2,6 @@
 import math
 import sys
 
-from .reference import TestbookObjectReference
-
 
 class Translator(object):
     @classmethod
@@ -74,7 +72,7 @@ def translate(cls, val):
             return cls.translate_list(val)
         elif isinstance(val, tuple):
             return cls.translate_tuple(val)
-        elif isinstance(val, TestbookObjectReference):
+        elif val.__class__.__name__ == "TestbookObjectReference":
             return val.name
 
         # Use this generic translation as a last resort