Refactor ipynb stuff into a new module, minor lint fixes

Author: ma-sadeghi

src/pytest_split/ipynb_compatibility.py

@@ -0,0 +1,60 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import List
+
+    from _pytest import nodes
+
+    from pytest_split.algorithms import TestGroup
+
+
+def _reorganize_broken_up_ipynbs(group: "TestGroup", items: "List[nodes.Item]") -> None:
+    """
+    Ensures that group doesn't contain partial IPy notebook cells.
+
+    ``pytest-split`` might, in principle, break up the cells of an
+    IPython notebook into different test groups, in which case the tests
+    most likely fail (for starters, libraries are imported in Cell 0, so
+    all subsequent calls to the imported libraries in the following cells
+    will raise ``NameError``).
+
+    """
+    item_node_ids = [item.nodeid for item in items]
+
+    # Deal with broken up notebooks at the beginning of the test group
+    if not group.selected or not _is_ipy_notebook(group.selected[0].nodeid):
+        return
+    first = group.selected[0].nodeid
+    siblings = _find_sibiling_ipynb_cells(first, item_node_ids)
+    if first != siblings[0]:
+        for item in list(group.selected):
+            if item.nodeid in siblings:
+                group.deselected.append(item)
+                group.selected.remove(item)
+
+    # Deal with broken up notebooks at the end of the test group
+    if not group.selected or not _is_ipy_notebook(group.selected[-1].nodeid):
+        return
+    last = group.selected[-1].nodeid
+    siblings = _find_sibiling_ipynb_cells(last, item_node_ids)
+    if last != siblings[-1]:
+        for item in list(group.deselected):
+            if item.nodeid in siblings:
+                group.deselected.remove(item)
+                group.selected.append(item)
+
+
+def _find_sibiling_ipynb_cells(ipynb_node_id: str, item_node_ids: "List[str]") -> "List[str]":
+    """
+    Returns all sibiling IPyNb cells given an IPyNb cell nodeid.
+    """
+    fpath = ipynb_node_id.split("::")[0]
+    return [item for item in item_node_ids if fpath in item]
+
+
+def _is_ipy_notebook(node_id: str) -> bool:
+    """
+    Returns True if node_id is an IPython notebook, otherwise False.
+    """
+    fpath = node_id.split("::")[0]
+    return fpath.endswith(".ipynb")

src/pytest_split/plugin.py

@@ -7,6 +7,7 @@
 from _pytest.reports import TestReport
 
 from pytest_split import algorithms
+from pytest_split.ipynb_compatibility import _reorganize_broken_up_ipynbs
 
 if TYPE_CHECKING:
     from typing import Dict, List, Optional, Union
@@ -16,7 +17,6 @@
     from _pytest.config.argparsing import Parser
     from _pytest.main import ExitCode
 
-    from pytest_split.algorithms import TestGroup
 
 # Ugly hack for freezegun compatibility: https://github.com/spulec/freezegun/issues/286
 STORE_DURATIONS_SETUP_AND_TEARDOWN_THRESHOLD = 60 * 10  # seconds
@@ -185,60 +185,6 @@ def pytest_collection_modifyitems(
         return None
 
 
-def _reorganize_broken_up_ipynbs(group: "TestGroup", items: list) -> None:
-    """
-    Ensures that group doesn't contain partial IPy notebook cells.
-
-    ``pytest-split`` might, in principle, break up the cells of an
-    IPython notebook into different test groups, in which case the tests
-    most likely fail (for starters, libraries are imported in Cell 0, so
-    all subsequent calls to the imported libraries in the following cells
-    will raise ``NameError``).
-
-    """
-    item_node_ids = [item.nodeid for item in items]
-
-    # Deal with broken up notebooks at the beginning of the test group
-    if not group.selected or not _is_ipy_notebook(group.selected[0].nodeid):
-        return
-    first = group.selected[0].nodeid
-    sibilings = _find_sibiling_ipynb_cells(first, item_node_ids)
-    if first != sibilings[0]:
-        for item in list(group.selected):
-            if item.nodeid in sibilings:
-                group.deselected.append(item)
-                group.selected.remove(item)
-
-    # Deal with broken up notebooks at the end of the test group
-    if not group.selected or not _is_ipy_notebook(group.selected[-1].nodeid):
-        return
-    last = group.selected[-1].nodeid
-    sibilings = _find_sibiling_ipynb_cells(last, item_node_ids)
-    if last != sibilings[-1]:
-        for item in list(group.deselected):
-            if item.nodeid in sibilings:
-                group.deselected.remove(item)
-                group.selected.append(item)
-
-
-def _find_sibiling_ipynb_cells(
-    ipynb_node_id: str, item_node_ids: "List[str]"
-) -> "List[str]":
-    """
-    Returns all sibiling IPyNb cells given an IPyNb cell nodeid.
-    """
-    fpath = ipynb_node_id.split("::")[0]
-    return [item for item in item_node_ids if fpath in item]
-
-
-def _is_ipy_notebook(node_id: str) -> bool:
-    """
-    Returns True if node_id is an IPython notebook, otherwise False.
-    """
-    fpath = node_id.split("::")[0]
-    return fpath.endswith(".ipynb")
-
-
 class PytestSplitCachePlugin(Base):
     """
     The cache plugin writes durations to our durations file.

tests/test_ipynb.py

@@ -3,7 +3,7 @@
 import pytest
 
 from pytest_split.algorithms import Algorithms
-from pytest_split.plugin import _reorganize_broken_up_ipynbs
+from pytest_split.ipynb_compatibility import _reorganize_broken_up_ipynbs
 
 item = namedtuple("item", "nodeid")
 
@@ -30,22 +30,22 @@ def test__reorganize_broken_up_ipynbs(self, algo_name):
         }
         items = [item(x) for x in durations.keys()]
         algo = Algorithms[algo_name].value
-        g1, g2, g3 = algo(splits=3, items=items, durations=durations)
+        groups = algo(splits=3, items=items, durations=durations)
 
-        assert g1.selected == [
+        assert groups[0].selected == [
             item(nodeid="temp/nbs/test_1.ipynb::Cell 0"),
             item(nodeid="temp/nbs/test_1.ipynb::Cell 1"),
             item(nodeid="temp/nbs/test_1.ipynb::Cell 2"),
             item(nodeid="temp/nbs/test_2.ipynb::Cell 0"),
             item(nodeid="temp/nbs/test_2.ipynb::Cell 1"),
         ]
-        assert g2.selected == [
+        assert groups[1].selected == [
             item(nodeid="temp/nbs/test_2.ipynb::Cell 2"),
             item(nodeid="temp/nbs/test_2.ipynb::Cell 3"),
             item(nodeid="temp/nbs/test_3.ipynb::Cell 0"),
             item(nodeid="temp/nbs/test_3.ipynb::Cell 1"),
         ]
-        assert g3.selected == [
+        assert groups[2].selected == [
             item(nodeid="temp/nbs/test_3.ipynb::Cell 2"),
             item(nodeid="temp/nbs/test_3.ipynb::Cell 3"),
             item(nodeid="temp/nbs/test_3.ipynb::Cell 4"),
@@ -54,8 +54,8 @@ def test__reorganize_broken_up_ipynbs(self, algo_name):
             item(nodeid="temp/nbs/test_4.ipynb::Cell 2"),
         ]
 
-        _reorganize_broken_up_ipynbs(g1, items)
-        assert g1.selected == [
+        _reorganize_broken_up_ipynbs(groups[0], items)
+        assert groups[0].selected == [
             item(nodeid="temp/nbs/test_1.ipynb::Cell 0"),
             item(nodeid="temp/nbs/test_1.ipynb::Cell 1"),
             item(nodeid="temp/nbs/test_1.ipynb::Cell 2"),
@@ -65,29 +65,18 @@ def test__reorganize_broken_up_ipynbs(self, algo_name):
             item(nodeid="temp/nbs/test_2.ipynb::Cell 3"),
         ]
 
-        _reorganize_broken_up_ipynbs(g2, items)
-        assert g2.selected == [
+        _reorganize_broken_up_ipynbs(groups[1], items)
+        assert groups[1].selected == [
             item(nodeid="temp/nbs/test_3.ipynb::Cell 0"),
             item(nodeid="temp/nbs/test_3.ipynb::Cell 1"),
             item(nodeid="temp/nbs/test_3.ipynb::Cell 2"),
             item(nodeid="temp/nbs/test_3.ipynb::Cell 3"),
             item(nodeid="temp/nbs/test_3.ipynb::Cell 4"),
         ]
 
-        _reorganize_broken_up_ipynbs(g3, items)
-        assert g3.selected == [
+        _reorganize_broken_up_ipynbs(groups[2], items)
+        assert groups[2].selected == [
             item(nodeid="temp/nbs/test_4.ipynb::Cell 0"),
             item(nodeid="temp/nbs/test_4.ipynb::Cell 1"),
             item(nodeid="temp/nbs/test_4.ipynb::Cell 2"),
         ]
-
-
-if __name__ == "__main__":
-    import warnings
-
-    warnings.filterwarnings("ignore")
-
-    import pytest
-
-    args = "test_ipynb.py"
-    retcode = pytest.main(args.split())