types: add types to dvc/tree (iterative#4947)

dvc/tree/base.py

@@ -2,6 +2,7 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from functools import partial
 from multiprocessing import cpu_count
+from typing import Any, ClassVar, Dict, Optional
 from urllib.parse import urlparse
 
 from funcy import cached_property, decorator
@@ -46,8 +47,8 @@ def use_state(call):
 
 class BaseTree:
     scheme = "base"
-    REQUIRES = {}
-    PATH_CLS = URLInfo
+    REQUIRES: ClassVar[Dict[str, str]] = {}
+    PATH_CLS = URLInfo  # type: Any
     JOBS = 4 * cpu_count()
 
     CHECKSUM_DIR_SUFFIX = ".dir"
@@ -59,9 +60,9 @@ class BaseTree:
     TRAVERSE_THRESHOLD_SIZE = 500000
     CAN_TRAVERSE = True
 
-    CACHE_MODE = None
+    CACHE_MODE: Optional[int] = None
     SHARED_MODE_MAP = {None: (None, None), "group": (None, None)}
-    PARAM_CHECKSUM = None
+    PARAM_CHECKSUM: ClassVar[Optional[str]] = None
 
     state = StateNoop()
 
@@ -155,9 +156,10 @@ def dir_mode(self):
     def cache(self):
         return getattr(self.repo.cache, self.scheme)
 
-    def open(self, path_info, mode="r", encoding=None):
+    def open(self, path_info, mode: str = "r", encoding: str = None):
         if hasattr(self, "_generate_download_url"):
-            func = self._generate_download_url  # noqa,pylint:disable=no-member
+            # pylint:disable=no-member
+            func = self._generate_download_url  # type: ignore[attr-defined]
             get_url = partial(func, path_info)
             return open_url(get_url, mode=mode, encoding=encoding)
 

dvc/tree/dvc.py

@@ -32,7 +32,6 @@ class DvcTree(BaseTree):  # pylint:disable=abstract-method
 
     scheme = "local"
     PARAM_CHECKSUM = "md5"
-    _dir_entry_hashes = {}
 
     def __init__(self, repo, fetch=False, stream=False):
         super().__init__(repo, {"url": repo.root_dir})

dvc/tree/local.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import stat
+from typing import Any, Dict
 
 from funcy import cached_property
 from shortuuid import uuid
@@ -27,7 +28,10 @@ class LocalTree(BaseTree):
     TRAVERSE_PREFIX_LEN = 2
 
     CACHE_MODE = 0o444
-    SHARED_MODE_MAP = {None: (0o644, 0o755), "group": (0o664, 0o775)}
+    SHARED_MODE_MAP: Dict[Any, Any] = {
+        None: (0o644, 0o755),
+        "group": (0o664, 0o775),
+    }
 
     def __init__(self, repo, config, use_dvcignore=False, dvcignore_root=None):
         super().__init__(repo, config)

dvc/tree/repo.py

@@ -4,7 +4,7 @@
 import threading
 from contextlib import suppress
 from itertools import takewhile
-from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Callable, Optional, Tuple, Type, Union
 
 from funcy import lfilter, wrap_with
 
@@ -22,12 +22,10 @@
 if TYPE_CHECKING:
     from dvc.repo import Repo
 
-    from .git import GitTree
-    from .local import LocalTree
-
-
 logger = logging.getLogger(__name__)
 
+RepoFactory = Union[Callable[[str], "Repo"], Type["Repo"]]
+
 
 class RepoTree(BaseTree):  # pylint:disable=abstract-method
     """DVC + git-tracked files tree.
@@ -43,18 +41,14 @@ class RepoTree(BaseTree):  # pylint:disable=abstract-method
     PARAM_CHECKSUM = "md5"
 
     def __init__(
-        self,
-        repo,
-        subrepos=False,
-        repo_factory: Callable[[str], "Repo"] = None,
-        **kwargs
+        self, repo, subrepos=False, repo_factory: RepoFactory = None, **kwargs
     ):
         super().__init__(repo, {"url": repo.root_dir})
 
         if not repo_factory:
             from dvc.repo import Repo
 
-            self.repo_factory = Repo
+            self.repo_factory: RepoFactory = Repo
         else:
             self.repo_factory = repo_factory
 
@@ -120,9 +114,7 @@ def _is_dvc_repo(self, dir_path):
         # dvcignore will ignore subrepos, therefore using `use_dvcignore=False`
         return self._main_repo.tree.isdir(repo_path, use_dvcignore=False)
 
-    def _get_tree_pair(
-        self, path
-    ) -> Tuple[Union["GitTree", "LocalTree"], DvcTree]:
+    def _get_tree_pair(self, path) -> Tuple[BaseTree, Optional[DvcTree]]:
         """
         Returns a pair of trees based on repo the path falls in, using prefix.
         """