Merge pull request #13 from datarootsio/avoid-typing-cast

Avoid typing cast

README.md

@@ -3,7 +3,7 @@
 # databooks
 [![maintained by dataroots](https://dataroots.io/maintained-rnd.svg)](https://dataroots.io)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
-[![Checked with mypy](http://www.mypy-lang.org/static/mypy_badge.svg)](http://mypy-lang.org/)
+[![Checked with mypy](https://img.shields.io/badge/mypy-checked-1f5082.svg)](http://mypy-lang.org/)
 [![Codecov](https://codecov.io/github/datarootsio/databooks/main/graph/badge.svg)](https://github.com/datarootsio/databooks/actions)
 [![tests](https://github.com/datarootsio/databooks/actions/workflows/test.yml/badge.svg)](https://github.com/datarootsio/databooks/actions)
 [![Downloads](https://pepy.tech/badge/databooks)](https://pepy.tech/project/databooks)

databooks/conflicts.py

@@ -3,28 +3,18 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, cast
+from typing import Any, Callable, List, Optional, Sequence, Tuple
 
 from git import Repo
 
 from databooks.common import write_notebook
-from databooks.data_models.base import BaseCells, DiffModel
 from databooks.data_models.notebook import Cell, Cells, JupyterNotebook
 from databooks.git_utils import ConflictFile, get_conflict_blobs, get_repo
 from databooks.logging import get_logger, set_verbose
 
 logger = get_logger(__file__)
 
 
-class DiffJupyterNotebook(DiffModel):
-    """Protocol for mypy static type checking."""
-
-    nbformat: int
-    nbformat_minor: int
-    metadata: Dict[str, Any]
-    cells: BaseCells[Any]
-
-
 def path2conflicts(
     nb_paths: List[Path], repo: Optional[Repo] = None
 ) -> List[ConflictFile]:
@@ -85,24 +75,17 @@ def conflict2nb(
         )
         logger.debug(msg)
 
-    if cell_fields_ignore:
-        for cells in (nb_1.cells, nb_2.cells):
-            for cell in cells:
-                cell.clear_metadata(
-                    cell_metadata_remove=[], cell_remove_fields=cell_fields_ignore
-                )
-
-    diff_nb = cast(DiffModel, nb_1 - nb_2)
-    nb = cast(
-        JupyterNotebook,
-        diff_nb.resolve(
-            ignore_none=ignore_none,
-            keep_first=meta_first,
-            keep_first_cells=cells_first,
-            first_id=conflict_file.first_log,
-            last_id=conflict_file.last_log,
-        ),
+    diff_nb = nb_1 - nb_2
+    nb = diff_nb.resolve(
+        ignore_none=ignore_none,
+        keep_first=meta_first,
+        keep_first_cells=cells_first,
+        first_id=conflict_file.first_log,
+        last_id=conflict_file.last_log,
     )
+    if not isinstance(nb, JupyterNotebook):
+        raise RuntimeError(f"Expected `databooks.JupyterNotebook`, got {type(nb)}.")
+
     logger.debug(f"Resolved conflicts in {conflict_file.filename}.")
     return nb
 

databooks/data_models/base.py

@@ -83,7 +83,7 @@ def resolve(
     if not is_diff:
         raise TypeError("Can only resolve dynamic 'diff models' (when `is_diff=True`).")
 
-    res_vals = cast(Dict[str, Any], {})
+    res_vals: Dict[str, Any] = {}
     for name, value in field_d.items():
         if isinstance(value, (DiffModel, BaseCells)):
             res_vals[name] = value.resolve(
@@ -134,7 +134,7 @@ def __str__(self) -> str:
         """Return outputs of __repr__."""
         return repr(self)
 
-    def __sub__(self, other: DatabooksBase) -> DatabooksBase:
+    def __sub__(self, other: DatabooksBase) -> DiffModel:
         """
         Subtraction between `databooks.data_models.base.DatabooksBase` objects.
 
@@ -152,7 +152,7 @@ def __sub__(self, other: DatabooksBase) -> DatabooksBase:
         other_d = dict(other)
 
         # Build dict with {field: (type, value)} for each field
-        fields_d = {}
+        fields_d: Dict[str, Any] = {}
         for name in self_d.keys() | other_d.keys():
             self_val = self_d.get(name)
             other_val = other_d.get(name)
@@ -166,11 +166,11 @@ def __sub__(self, other: DatabooksBase) -> DatabooksBase:
                 fields_d[name] = (tuple, (self_val, other_val))
 
         # Build Pydantic models dynamically
-        DiffModel = create_model(
+        DiffInstance = create_model(
             "Diff" + type(self).__name__,
             __base__=type(self),
             resolve=resolve,
             is_diff=True,
-            **cast(Dict[str, Any], fields_d),
+            **fields_d,
         )
-        return DiffModel()  # it'll be filled in with the defaults
+        return cast(DiffModel, DiffInstance())  # it'll be filled in with the defaults

databooks/git_utils.py

@@ -1,7 +1,7 @@
 """Git helper functions."""
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, cast
+from typing import Dict, List
 
 from git import Blob, Git, Repo  # type: ignore
 
@@ -55,8 +55,12 @@ def get_conflict_blobs(repo: Repo) -> List[ConflictFile]:
         for k, v in unmerged_blobs.items()
         if 0 not in dict(v).keys()  # only get blobs that could not be merged
     )
-    # Type checking: `repo.working_dir` is not None if repo is passed
-    repo.working_dir = cast(str, repo.working_dir)
+
+    if not isinstance(repo.working_dir, (Path, str)):
+        raise RuntimeError(
+            "Expected `repo` to be `pathlib.Path` or `str`, got"
+            f" {type(repo.working_dir)}."
+        )
     return [
         ConflictFile(
             filename=repo.working_dir / blob.filename,

tests/test_conflicts.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-from typing import cast
 
 from py._path.local import LocalPath
 
@@ -38,7 +37,8 @@ def test_path2diff(tmpdir: LocalPath) -> None:
         commit_message_main="Commit message from main",
         commit_message_other="Commit message from other",
     )
-    git_repo.working_dir = cast(Path, git_repo.working_dir)
+
+    assert isinstance(git_repo.working_dir, (Path, str))
 
     conflict_files = path2conflicts(nb_paths=[nb_filepath], repo=git_repo)
 

tests/test_data_models/test_base.py

@@ -1,13 +1,11 @@
-from typing import cast
-
-from databooks.data_models.base import DatabooksBase, DiffModel
+from databooks.data_models.base import DatabooksBase
 
 
 def test_base_sub() -> None:
     """Use the `-` operator and resolve the diffs from the base model."""
     model_1 = DatabooksBase(test=0, foo=1, bar="2")
     model_2 = DatabooksBase(bar=4, foo=2, baz=2.3, test=0)
-    diff = cast(DiffModel, model_1 - model_2)
+    diff = model_1 - model_2
     assert diff.__class__.__name__ == "DiffDatabooksBase"
     assert dict(diff) == {
         "is_diff": True,

tests/test_data_models/test_notebook.py

@@ -2,12 +2,11 @@
 import logging
 from copy import deepcopy
 from importlib import resources
-from typing import List, Tuple, cast
+from typing import List, Tuple
 
 import pytest
 from _pytest.logging import LogCaptureFixture
 
-from databooks.data_models.base import DiffModel
 from databooks.data_models.notebook import (
     Cell,
     CellMetadata,
@@ -203,7 +202,7 @@ def test_notebook_sub(self) -> None:
         )
         notebook_2.cells = notebook_2.cells + [extra_cell]
 
-        diff = cast(DiffModel, notebook_1 - notebook_2)
+        diff = notebook_1 - notebook_2
         notebook = deepcopy(notebook_1)
 
         # add `tags` since we resolve with default `ignore_none = True`

tests/test_git_utils.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-from typing import cast
 
 from git import GitCommandError, Repo
 from py._path.local import LocalPath
@@ -18,7 +17,12 @@ def init_repo_conflicts(
 ) -> Repo:
     """Create git repo and create a conflict."""
     git_repo = Repo.init(path=tmpdir)
-    git_repo.working_dir = cast(Path, git_repo.working_dir)
+
+    if not isinstance(git_repo.working_dir, (Path, str)):
+        raise RuntimeError(
+            f"Expected `pathlib.Path` or `str`, got {type(git_repo.working_dir)}."
+        )
+
     git_filepath = git_repo.working_dir / filename
 
     git_repo.git.checkout("-b", "main")
@@ -64,7 +68,9 @@ def test_get_conflict_blobs(tmpdir: LocalPath) -> None:
         commit_message_main="Commit message from main",
         commit_message_other="Commit message from other",
     )
-    git_repo.working_dir = cast(Path, git_repo.working_dir)
+
+    assert isinstance(git_repo.working_dir, (Path, str))
+
     conflicts = get_conflict_blobs(repo=git_repo)
     assert len(conflicts) == 1