Fix typechecking errors - 2 (#4946)

dvc/command/dag.py

@@ -2,6 +2,7 @@ import argparse
 import logging
 
 from dvc.command.base import CmdBase, append_doc_link
+from dvc.stage import Stage
 
 logger = logging.getLogger(__name__)
 
@@ -52,18 +53,10 @@ def _collect_targets(repo, target, outs):
 def _transform(repo, outs):
     import networkx as nx
 
-    if outs:
-        G = repo.outs_graph
-
-        def _relabel(out):
-            return str(out)
-
-    else:
-        G = repo.graph
-
-        def _relabel(stage):
-            return stage.addressing
+    def _relabel(node) -> str:
+        return node.addressing if isinstance(node, Stage) else str(node)
 
+    G = repo.outs_graph if outs else repo.graph
     return nx.relabel_nodes(G, _relabel, copy=True)
 
 

dvc/command/experiments.py

@@ -51,13 +51,13 @@ def _filter_names(
                 update_func({match: None for match in matches})
 
     if include:
-        ret = OrderedDict()
+        ret: OrderedDict = OrderedDict()
         _filter(include, ret.update)
     else:
         ret = OrderedDict({name: None for name in names})
 
     if exclude:
-        _filter(exclude, ret.difference_update)
+        _filter(exclude, ret.difference_update)  # type: ignore[attr-defined]
 
     return [".".join(name) for name in ret]
 

dvc/dependency/base.py

@@ -1,3 +1,5 @@
+from typing import Type
+
 from dvc.exceptions import DvcException
 
 
@@ -21,9 +23,11 @@ class DependencyIsStageFileError(DvcException):
 class BaseDependency:
     IS_DEPENDENCY = True
 
-    DoesNotExistError = DependencyDoesNotExistError
-    IsNotFileOrDirError = DependencyIsNotFileOrDirError
-    IsStageFileError = DependencyIsStageFileError
+    DoesNotExistError = DependencyDoesNotExistError  # type: Type[DvcException]
+    IsNotFileOrDirError = (
+        DependencyIsNotFileOrDirError
+    )  # type: Type[DvcException]
+    IsStageFileError = DependencyIsStageFileError  # type: Type[DvcException]
 
     def update(self, rev=None):
         pass

dvc/output/base.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from copy import copy
+from typing import Type
 from urllib.parse import urlparse
 
 from voluptuous import Any
@@ -85,10 +86,10 @@ class BaseOutput:
         },
     )
 
-    DoesNotExistError = OutputDoesNotExistError
-    IsNotFileOrDirError = OutputIsNotFileOrDirError
-    IsStageFileError = OutputIsStageFileError
-    IsIgnoredError = OutputIsIgnoredError
+    DoesNotExistError = OutputDoesNotExistError  # type: Type[DvcException]
+    IsNotFileOrDirError = OutputIsNotFileOrDirError  # type: Type[DvcException]
+    IsStageFileError = OutputIsStageFileError  # type: Type[DvcException]
+    IsIgnoredError = OutputIsIgnoredError  # type: Type[DvcException]
 
     sep = "/"
 

dvc/output/gs.py

@@ -1,7 +1,7 @@
-from dvc.output.s3 import S3Output
+from dvc.output.base import BaseOutput
 
 from ..tree.gs import GSTree
 
 
-class GSOutput(S3Output):
+class GSOutput(BaseOutput):
     TREE_CLS = GSTree

dvc/parsing/__init__.py

@@ -159,7 +159,7 @@ class DataResolver:
             stage_name=name,
         )
 
-        logger.trace(  # pytype: disable=attribute-error
+        logger.trace(  # type: ignore[attr-defined]
             "Context during resolution of stage %s:\n%s", name, context
         )
 

dvc/parsing/context.py

@@ -190,7 +190,7 @@ class Container(Node, ABC):
         index = index.strip()
         index = self._key_transform(index)
         try:
-            d = self.data[index]
+            d = self[index]
         except LookupError as exc:
             raise ValueError(
                 f"Could not find '{index}' in {self.data}"

dvc/remote/local.py

@@ -3,4 +3,4 @@ from .index import RemoteIndexNoop
 
 
 class LocalRemote(Remote):
-    INDEX_CLS = RemoteIndexNoop
+    INDEX_CLS = RemoteIndexNoop  # type: ignore[assignment]

dvc/remote/slow_link_detection.py

@@ -1,3 +1,4 @@
+# type: ignore
 import logging
 import sys
 import time

dvc/repo/params/show.py

@@ -1,4 +1,5 @@
 import logging
+from typing import TYPE_CHECKING, List
 
 from dvc.dependency.param import ParamsDependency
 from dvc.exceptions import DvcException
@@ -7,6 +8,10 @@ from dvc.repo import locked
 from dvc.repo.collect import collect
 from dvc.utils.serialize import LOADERS, ParseError
 
+if TYPE_CHECKING:
+    from dvc.output.base import BaseOutput
+    from dvc.repo import Repo
+
 logger = logging.getLogger(__name__)
 
 
@@ -14,11 +19,11 @@ class NoParamsError(DvcException):
     pass
 
 
-def _is_params(dep):
+def _is_params(dep: "BaseOutput"):
     return isinstance(dep, ParamsDependency)
 
 
-def _collect_configs(repo, rev):
+def _collect_configs(repo: "Repo", rev) -> List[PathInfo]:
     params, _ = collect(repo, deps=True, output_filter=_is_params, rev=rev)
     configs = {p.path_info for p in params}
     configs.add(PathInfo(repo.root_dir) / ParamsDependency.DEFAULT_PARAMS_FILE)

dvc/repo/plots/template.py

@@ -1,5 +1,6 @@
 import json
 import os
+from typing import Any, Dict, Optional
 
 from funcy import cached_property
 
@@ -29,8 +30,8 @@ class Template:
     EXTENSION = ".json"
     ANCHOR = "<DVC_METRIC_{}>"
 
-    DEFAULT_CONTENT = None
-    DEFAULT_NAME = None
+    DEFAULT_CONTENT: Optional[Dict[str, Any]] = None
+    DEFAULT_NAME: Optional[str] = None
 
     def __init__(self, content=None, name=None):
         if content: