add testing (#4)

* remove scratch notebook

* ignore vscode/pytest_cache

* initial updates for testing

* add docstring back

* testing might be interfering too much here

* pass optional ipython shell arg to register/deregister

* remove these other fixtures since we don't need them

* revert register()/deregister() tests since they use the ipython shell instance in a better way

.gitignore

@@ -1,4 +1,7 @@
 __pycache__/
 *.pyc
 
-dist/
+dist/
+
+.pytest_cache
+.vscode

dx/formatters.py

@@ -1,9 +1,16 @@
 import pandas as pd
 from IPython import get_ipython
+from IPython.core.interactiveshell import InteractiveShell
 from IPython.core.formatters import DisplayFormatter
 from pandas.io.json import build_table_schema
+from typing import Optional
+
+IN_IPYTHON_ENV = get_ipython() is not None
+
+DEFAULT_IPYTHON_DISPLAY_FORMATTER = None
+if IN_IPYTHON_ENV:
+    DEFAULT_IPYTHON_DISPLAY_FORMATTER = get_ipython().display_formatter
 
-DEFAULT_IPYTHON_DISPLAY_FORMATTER = get_ipython().display_formatter
 DX_MEDIA_TYPE = "application/vnd.dex.v1+json"
 
 
@@ -32,16 +39,27 @@ def format_dx(df) -> tuple:
     return (payload, metadata)
 
 
-def deregister() -> None:
+def deregister(ipython_shell: Optional[InteractiveShell] = None) -> None:
     """Reverts IPython.display_formatter to its original state"""
+    if not IN_IPYTHON_ENV and ipython_shell is None:
+        return
     pd.options.display.max_rows = 60
-    get_ipython().display_formatter = DEFAULT_IPYTHON_DISPLAY_FORMATTER
+    ipython = ipython_shell or get_ipython()
+    ipython.display_formatter = DEFAULT_IPYTHON_DISPLAY_FORMATTER
 
 
-def register() -> None:
+def register(ipython_shell: Optional[InteractiveShell] = None) -> None:
     """Overrides the default IPython display formatter to use DXDisplayFormatter"""
+    if not IN_IPYTHON_ENV and ipython_shell is None:
+        return
     pd.options.display.max_rows = 100_000
-    get_ipython().display_formatter = DXDisplayFormatter()
+
+    if ipython_shell is not None:
+        global DEFAULT_IPYTHON_DISPLAY_FORMATTER
+        DEFAULT_IPYTHON_DISPLAY_FORMATTER = ipython_shell.display_formatter
+
+    ipython = ipython_shell or get_ipython()
+    ipython.display_formatter = DXDisplayFormatter()
 
 
 disable = deregister

dx/tests/conftest.py

@@ -0,0 +1,54 @@
+import random
+import string
+
+import pandas as pd
+import pytest
+from IPython.core import formatters
+from IPython.core.formatters import BaseFormatter, DisplayFormatter
+from IPython.terminal.interactiveshell import TerminalInteractiveShell
+from IPython.testing import tools
+
+
+@pytest.fixture
+def get_ipython() -> TerminalInteractiveShell:
+    if TerminalInteractiveShell._instance:
+        return TerminalInteractiveShell.instance()
+
+    config = tools.default_config()
+    config.TerminalInteractiveShell.simple_prompt = True
+    shell = TerminalInteractiveShell.instance(config=config)
+    return shell
+
+
+@pytest.fixture
+def sample_dataframe() -> pd.DataFrame:
+    df = pd.DataFrame(
+        {
+            "col_1": list("aaa"),
+            "col_2": list("bbb"),
+            "col_3": list("ccc"),
+        }
+    )
+    return df
+
+
+@pytest.fixture
+def sample_large_dataframe() -> pd.DataFrame:
+    """
+    Generates a random assortment of values of different data types,
+    and returns a larger dataframe than the `sample_dataframe` fixture.
+    """
+    n_rows = 100_000
+    data = {
+        "float_col": [random.randint(0, 100) for _ in range(n_rows)],
+        "int_col": [random.random() for _ in range(n_rows)],
+        "bool_col": [random.choice([True, False]) for _ in range(n_rows)],
+        "str_col": [
+            "".join(random.sample(string.ascii_uppercase, 3)) for _ in range(n_rows)
+        ],
+        "date_col": [
+            pd.Timestamp("now") - pd.Timedelta(hours=random.randint(-100, 100))
+            for _ in range(n_rows)
+        ],
+    }
+    return pd.DataFrame(data)

dx/tests/test_formatters.py

@@ -0,0 +1,72 @@
+from dx.formatters import (
+    DX_MEDIA_TYPE,
+    DXDisplayFormatter,
+    deregister,
+    format_dx,
+    register,
+)
+from IPython.core.formatters import DisplayFormatter
+from IPython.terminal.interactiveshell import TerminalInteractiveShell
+
+
+def test_media_type(sample_dataframe):
+    payload, _ = format_dx(sample_dataframe)
+    assert DX_MEDIA_TYPE in payload
+
+
+def test_data_structure(sample_dataframe):
+    """
+    The transformed data needs to represent a list of lists,
+    each associated with a column in the dataframe,
+    including one for the dataframe's index.
+    """
+    payload, _ = format_dx(sample_dataframe)
+    data = payload[DX_MEDIA_TYPE]["data"]
+    assert isinstance(data, list)
+    assert len(data) == 4
+    assert isinstance(data[0], list)
+
+
+def test_data_list_order(sample_dataframe):
+    """
+    Ensure the payload contains lists as column values,
+    and not row values.
+    """
+    payload, _ = format_dx(sample_dataframe)
+    data = payload[DX_MEDIA_TYPE]["data"]
+    assert data[0] == [0, 1, 2]  # index values
+    assert data[1] == list("aaa")  # "col_1" values
+    assert data[2] == list("bbb")  # "col_2" values
+    assert data[3] == list("ccc")  # "col_3" values
+
+
+def test_fields_match_data_length(sample_dataframe):
+    """
+    The number of fields in the schema needs to match
+    the number of lists in the data list.
+    """
+    payload, _ = format_dx(sample_dataframe)
+    data = payload[DX_MEDIA_TYPE]["data"]
+    fields = payload[DX_MEDIA_TYPE]["schema"]["fields"]
+    assert len(data) == len(fields)
+
+
+def test_register_ipython_display_formatter(get_ipython: TerminalInteractiveShell):
+    """
+    Test that the display formatter for an IPython shell is
+    successfully registered as a DXDisplayFormatter.
+    """
+    register(ipython_shell=get_ipython)
+    assert isinstance(get_ipython.display_formatter, DXDisplayFormatter)
+
+
+def test_deregister_ipython_display_formatter(get_ipython: TerminalInteractiveShell):
+    """
+    Test that the display formatter reverts to the default
+    `IPython.core.formatters.DisplayFormatter` after deregistering.
+    """
+    register(ipython_shell=get_ipython)
+    assert isinstance(get_ipython.display_formatter, DXDisplayFormatter)
+
+    deregister(ipython_shell=get_ipython)
+    assert isinstance(get_ipython.display_formatter, DisplayFormatter)