From 04c53601c027c357326dc0f576501030a3cc8975 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Tue, 24 Nov 2020 11:16:28 +0545
Subject: [PATCH] types: add types to dvc/tree (#4947)

---
 dvc/tree/base.py  | 14 ++++++++------
 dvc/tree/dvc.py   |  1 -
 dvc/tree/local.py |  6 +++++-
 dvc/tree/repo.py  | 20 ++++++--------------
 4 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/dvc/tree/base.py b/dvc/tree/base.py
index 6d7fd8e94..0848fd8e1 100644
--- a/dvc/tree/base.py
+++ b/dvc/tree/base.py
@@ -2,6 +2,7 @@ import logging
 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 @@ class BaseTree:
     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)
 
diff --git a/dvc/tree/dvc.py b/dvc/tree/dvc.py
index 5d09c642a..893beeb77 100644
--- a/dvc/tree/dvc.py
+++ b/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})
diff --git a/dvc/tree/local.py b/dvc/tree/local.py
index dab597ec2..5eddad93e 100644
--- a/dvc/tree/local.py
+++ b/dvc/tree/local.py
@@ -2,6 +2,7 @@ import errno
 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)
diff --git a/dvc/tree/repo.py b/dvc/tree/repo.py
index e696013e5..65c9a5e27 100644
--- a/dvc/tree/repo.py
+++ b/dvc/tree/repo.py
@@ -4,7 +4,7 @@ import stat
 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 @@ from .dvc import DvcTree
 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 @@ class RepoTree(BaseTree):  # pylint:disable=abstract-method
         # 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.
         """