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