Add misc type hints (#8430)

dvc/cli/command.py

@@ -1,6 +1,10 @@
 import logging
 import os
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from dvc.config import Config
 
 logger = logging.getLogger(__name__)
 
@@ -8,13 +12,13 @@ logger = logging.getLogger(__name__)
 class CmdBase(ABC):
     UNINITIALIZED = False
 
-    def __init__(self, args):
+    def __init__(self, args: Any):
         from dvc.repo import Repo
 
         os.chdir(args.cd)
 
         self.repo: "Repo" = Repo(uninitialized=self.UNINITIALIZED)
-        self.config = self.repo.config
+        self.config: "Config" = self.repo.config
         self.args = args
 
     def do_run(self):

dvc/commands/machine.py

@@ -1,4 +1,5 @@
 import argparse
+from typing import Dict, List
 
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link, fix_subparsers
@@ -6,7 +7,6 @@ from dvc.commands.config import CmdConfig
 from dvc.compare import TabularData
 from dvc.config import ConfigError
 from dvc.exceptions import DvcException
-from dvc.types import Dict, List
 from dvc.ui import ui
 from dvc.utils import format_link
 

dvc/commands/plots.py

@@ -94,7 +94,7 @@ class CmdPlots(CmdBase):
                 )
         return html_template_path
 
-    def run(self):  # noqa: C901
+    def run(self) -> int:  # noqa: C901
         from pathlib import Path
 
         from dvc.render.match import match_defs_renderers

dvc/commands/queue/status.py

@@ -1,6 +1,6 @@
 import argparse
 import logging
-from typing import List, Mapping, Optional
+from typing import Dict, List, Optional
 
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
@@ -14,19 +14,21 @@ logger = logging.getLogger(__name__)
 class CmdQueueStatus(CmdBase):
     """Show queue task and worker status."""
 
-    def run(self):
+    def run(self) -> int:
         result: List[
-            Mapping[str, Optional[str]]
+            Dict[str, Optional[str]]
         ] = self.repo.experiments.celery_queue.status()
         if result:
             all_headers = ["Task", "Name", "Created", "Status"]
             td = TabularData(all_headers)
             for exp in result:
                 created = format_time(exp.get("timestamp"))
+                assert exp["rev"]
+                assert exp["status"]
                 td.append(
                     [
                         exp["rev"][:7],
-                        exp.get("name", ""),
+                        exp.get("name") or "",
                         created,
                         exp["status"],
                     ]

dvc/config.py

@@ -4,14 +4,18 @@ import os
 import re
 from contextlib import contextmanager
 from functools import partial
-from typing import Dict
+from typing import TYPE_CHECKING, Dict, Optional
 
 from funcy import compact, memoize, re_find
 
 from dvc.exceptions import DvcException, NotDvcRepoError
+from dvc.types import DictStrAny, StrPath
 
 from .utils.objects import cached_property
 
+if TYPE_CHECKING:
+    from dvc.fs import FileSystem
+
 logger = logging.getLogger(__name__)
 
 
@@ -87,7 +91,11 @@ class Config(dict):
     CONFIG_LOCAL = "config.local"
 
     def __init__(
-        self, dvc_dir=None, validate=True, fs=None, config=None
+        self,
+        dvc_dir: Optional[StrPath] = None,
+        validate: bool = True,
+        fs: Optional["FileSystem"] = None,
+        config: Optional[DictStrAny] = None,
     ):  # pylint: disable=super-init-not-called
         from dvc.fs import LocalFileSystem
 
@@ -145,7 +153,7 @@ class Config(dict):
         with open(config_file, "w+", encoding="utf-8"):
             return Config(dvc_dir)
 
-    def load(self, validate=True, config=None):
+    def load(self, validate: bool = True, config: Optional[DictStrAny] = None):
         """Loads config from all the config files.
 
         Raises:

dvc/dependency/repo.py

@@ -1,7 +1,16 @@
 import os
 from collections import defaultdict
 from copy import copy
-from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple
+from typing import (
+    TYPE_CHECKING,
+    ContextManager,
+    Dict,
+    Iterable,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
 
 from voluptuous import Required
 
@@ -11,6 +20,9 @@ from dvc_data.hashfile.meta import Meta
 from .base import Dependency
 
 if TYPE_CHECKING:
+    from dvc.output import Output
+    from dvc.repo import Repo
+    from dvc.stage import Stage
     from dvc_data.hashfile.hash_info import HashInfo
     from dvc_data.hashfile.obj import HashFile
     from dvc_objects.db import ObjectDB
@@ -30,7 +42,9 @@ class RepoDependency(Dependency):
         }
     }
 
-    def __init__(self, def_repo, stage, *args, **kwargs):
+    def __init__(
+        self, def_repo: Dict[str, str], stage: "Stage", *args, **kwargs
+    ):
         self.def_repo = def_repo
         self._objs: Dict[str, "HashFile"] = {}
         self._meta: Dict[str, Meta] = {}
@@ -61,10 +75,10 @@ class RepoDependency(Dependency):
     def save(self):
         pass
 
-    def dumpd(self, **kwargs):
+    def dumpd(self, **kwargs) -> Dict[str, Union[str, Dict[str, str]]]:
         return {self.PARAM_PATH: self.def_path, self.PARAM_REPO: self.def_repo}
 
-    def download(self, to, jobs=None):
+    def download(self, to: "Output", jobs: Optional[int] = None):
         from dvc_data.hashfile.checkout import checkout
 
         for odb, objs in self.get_used_objs().items():
@@ -81,13 +95,13 @@ class RepoDependency(Dependency):
             prompt=confirm,
         )
 
-    def update(self, rev=None):
+    def update(self, rev: Optional[str] = None):
         if rev:
             self.def_repo[self.PARAM_REV] = rev
         with self._make_repo(locked=False) as repo:
             self.def_repo[self.PARAM_REV_LOCK] = repo.get_rev()
 
-    def changed_checksum(self):
+    def changed_checksum(self) -> bool:
         # From current repo point of view what describes RepoDependency is its
         # origin project url and rev_lock, and it makes RepoDependency
         # immutable, hence its impossible for checksum to change.
@@ -100,7 +114,7 @@ class RepoDependency(Dependency):
         return used
 
     def _get_used_and_obj(
-        self, obj_only=False, **kwargs
+        self, obj_only: bool = False, **kwargs
     ) -> Tuple[Dict[Optional["ObjectDB"], Set["HashInfo"]], Meta, "HashFile"]:
         from dvc.config import NoRemoteError
         from dvc.exceptions import NoOutputOrStageError, PathMissingError
@@ -154,7 +168,9 @@ class RepoDependency(Dependency):
                 used_obj_ids[object_store].update(oid for _, _, oid in obj)
             return used_obj_ids, meta, obj
 
-    def _check_circular_import(self, odb, obj_ids):
+    def _check_circular_import(
+        self, odb: "ObjectDB", obj_ids: Iterable[str]
+    ) -> None:
         from dvc.exceptions import CircularImportError
         from dvc.fs.dvc import DVCFileSystem
         from dvc_data.hashfile.db.reference import ReferenceHashFileDB
@@ -209,14 +225,16 @@ class RepoDependency(Dependency):
         )
         return meta
 
-    def _make_repo(self, locked=True, **kwargs):
+    def _make_repo(
+        self, locked: bool = True, **kwargs
+    ) -> "ContextManager[Repo]":
         from dvc.external_repo import external_repo
 
         d = self.def_repo
         rev = self._get_rev(locked=locked)
         return external_repo(d[self.PARAM_URL], rev=rev, **kwargs)
 
-    def _get_rev(self, locked=True):
+    def _get_rev(self, locked: bool = True):
         d = self.def_repo
         return (d.get(self.PARAM_REV_LOCK) if locked else None) or d.get(
             self.PARAM_REV

dvc/dvcfile.py

@@ -31,6 +31,7 @@ if TYPE_CHECKING:
     from dvc.repo import Repo
 
     from .parsing import DataResolver
+    from .stage import Stage
 
 logger = logging.getLogger(__name__)
 _T = TypeVar("_T")
@@ -198,16 +199,16 @@ class SingleStageFile(FileMixin):
     params: List[str] = []
 
     @property
-    def stage(self):
+    def stage(self) -> "Stage":
         data, raw = self._load()
         return self.LOADER.load_stage(self, data, raw)
 
     @property
-    def stages(self):
+    def stages(self) -> LOADER:
         data, raw = self._load()
         return self.LOADER(self, data, raw)
 
-    def dump(self, stage, **kwargs):
+    def dump(self, stage, **kwargs) -> None:
         """Dumps given stage appropriately in the dvcfile."""
         from dvc.stage import PipelineStage
 

dvc/exceptions.py

@@ -1,4 +1,6 @@
 """Exceptions raised by the dvc."""
+from typing import Dict, List
+
 from dvc.utils import format_link
 
 
@@ -227,7 +229,7 @@ class UploadError(FileTransferError):
 
 
 class CheckoutError(DvcException):
-    def __init__(self, target_infos, stats=None):
+    def __init__(self, target_infos: List[str], stats: Dict[str, List[str]]):
         from dvc.utils import error_link
 
         self.target_infos = target_infos

dvc/external_repo.py

@@ -3,7 +3,7 @@ import os
 import tempfile
 import threading
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Collection, Dict, Iterator, Optional
 
 from funcy import retry, wrap_with
 
@@ -17,6 +17,7 @@ from dvc.exceptions import (
 )
 from dvc.repo import Repo
 from dvc.scm import CloneError, map_scm_exception
+from dvc.types import StrPath
 from dvc.utils import relpath
 
 if TYPE_CHECKING:
@@ -28,8 +29,13 @@ logger = logging.getLogger(__name__)
 @contextmanager
 @map_scm_exception()
 def external_repo(
-    url, rev=None, for_write=False, cache_dir=None, cache_types=None, **kwargs
-):
+    url,
+    rev: Optional[str] = None,
+    for_write: bool = False,
+    cache_dir: Optional[StrPath] = None,
+    cache_types: Optional[Collection[str]] = None,
+    **kwargs,
+) -> Iterator["Repo"]:
     from scmrepo.git import Git
 
     from dvc.config import NoRemoteError

dvc/fs/dvc.py

@@ -6,11 +6,21 @@ import os
 import posixpath
 import threading
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 from fsspec.spec import AbstractFileSystem
 from funcy import wrap_with
 
+from dvc.types import StrPath
 from dvc_objects.fs.base import FileSystem
 from dvc_objects.fs.path import Path
 
@@ -155,7 +165,7 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
         if hasattr(repo, "dvc_dir"):
             self._datafss[key] = DataFileSystem(index=repo.index.data["repo"])
 
-    def _get_key(self, path) -> Key:
+    def _get_key(self, path: StrPath) -> Key:
         parts = self.repo.fs.path.relparts(path, self.repo.root_dir)
         if parts == (os.curdir,):
             return ()
@@ -372,29 +382,30 @@ class DVCFileSystem(FileSystem):
     protocol = "local"
     PARAM_CHECKSUM = "md5"
 
-    def _prepare_credentials(self, **config):
+    def _prepare_credentials(self, **config) -> Dict[str, Any]:
         return config
 
     @functools.cached_property
-    def fs(self):  # pylint: disable=invalid-overridden-method
+    # pylint: disable-next=invalid-overridden-method
+    def fs(self) -> "DVCFileSystem":
         return _DVCFileSystem(**self.fs_args)
 
-    def isdvc(self, path, **kwargs):
+    def isdvc(self, path, **kwargs) -> bool:
         return self.fs.isdvc(path, **kwargs)
 
     @property
-    def path(self):  # pylint: disable=invalid-overridden-method
+    def path(self) -> Path:  # pylint: disable=invalid-overridden-method
         return self.fs.path
 
     @property
-    def repo(self):
+    def repo(self) -> "Repo":
         return self.fs.repo
 
     @property
-    def repo_url(self):
+    def repo_url(self) -> str:
         return self.fs.repo_url
 
-    def from_os_path(self, path):
+    def from_os_path(self, path: str) -> str:
         if os.path.isabs(path):
             path = os.path.relpath(path, self.repo.root_dir)
 

dvc/ignore.py

@@ -3,6 +3,7 @@ import os
 import re
 from collections import namedtuple
 from itertools import chain, groupby, takewhile
+from typing import List, Optional
 
 from pathspec.patterns import GitWildMatchPattern
 from pathspec.util import normalize_file
@@ -10,7 +11,6 @@ from pygtrie import Trie
 
 from dvc.fs import AnyFSPath, FileSystem, Schemes, localfs
 from dvc.pathspec_math import PatternInfo, merge_patterns
-from dvc.types import List, Optional
 
 logger = logging.getLogger(__name__)
 

dvc/machine/__init__.py

@@ -127,6 +127,7 @@ class MachineManager:
         self, repo: "Repo", backends: Optional[Iterable[str]] = None, **kwargs
     ):
         self.repo = repo
+        assert self.repo.tmp_dir
         tmp_dir = os.path.join(self.repo.tmp_dir, "machine")
         self.backends = MachineBackends(backends, tmp_dir=tmp_dir, **kwargs)
 

dvc/output.py

@@ -32,7 +32,7 @@ from dvc_objects.errors import ObjectFormatError
 
 from .annotations import ANNOTATION_FIELDS, ANNOTATION_SCHEMA, Annotation
 from .fs import LocalFileSystem, RemoteMissingDepsError, Schemes, get_cloud_fs
-from .fs.callbacks import DEFAULT_CALLBACK
+from .fs.callbacks import DEFAULT_CALLBACK, Callback
 from .utils import relpath
 from .utils.fs import path_isin
 
@@ -41,7 +41,7 @@ if TYPE_CHECKING:
     from dvc_data.index import DataIndexKey
     from dvc_objects.db import ObjectDB
 
-    from .fs.callbacks import Callback
+    from .ignore import DvcIgnoreFilter
 
 logger = logging.getLogger(__name__)
 
@@ -561,7 +561,7 @@ class Output:
             return self.meta.version_id != self.get_meta().version_id
         return False
 
-    def workspace_status(self):
+    def workspace_status(self) -> Dict[str, str]:
         if not self.exists:
             return {str(self): "deleted"}
 
@@ -573,40 +573,40 @@ class Output:
 
         return {}
 
-    def status(self):
+    def status(self) -> Dict[str, str]:
         if self.hash_info and self.use_cache and self.changed_cache():
             return {str(self): "not in cache"}
 
         return self.workspace_status()
 
-    def changed(self):
+    def changed(self) -> bool:
         status = self.status()
         logger.debug(str(status))
         return bool(status)
 
     @property
-    def dvcignore(self):
+    def dvcignore(self) -> Optional["DvcIgnoreFilter"]:
         if self.fs.protocol == "local":
             return self.repo.dvcignore
         return None
 
     @property
-    def is_empty(self):
+    def is_empty(self) -> bool:
         return self.fs.is_empty(self.fs_path)
 
-    def isdir(self):
+    def isdir(self) -> bool:
         if self._is_path_dvcignore(self.fs_path):
             return False
         return self.fs.isdir(self.fs_path)
 
-    def isfile(self):
+    def isfile(self) -> bool:
         if self._is_path_dvcignore(self.fs_path):
             return False
         return self.fs.isfile(self.fs_path)
 
     # pylint: disable=no-member
 
-    def ignore(self):
+    def ignore(self) -> None:
         if not self.use_scm_ignore:
             return
 
@@ -615,7 +615,7 @@ class Output:
 
         self.repo.scm_context.ignore(self.fspath)
 
-    def ignore_remove(self):
+    def ignore_remove(self) -> None:
         if not self.use_scm_ignore:
             return
 
@@ -623,11 +623,11 @@ class Output:
 
     # pylint: enable=no-member
 
-    def save(self):
+    def save(self) -> None:
         if not self.exists:
             raise self.DoesNotExistError(self)
 
-        if not self.isfile and not self.isdir:
+        if not self.isfile() and not self.isdir():
             raise self.IsNotFileOrDirError(self)
 
         if self.is_empty:
@@ -663,11 +663,11 @@ class Output:
         self.hash_info = self.obj.hash_info
         self.files = None
 
-    def set_exec(self):
+    def set_exec(self) -> None:
         if self.isfile() and self.meta.isexec:
             self.odb.set_exec(self.fs_path)
 
-    def _checkout(self, *args, **kwargs):
+    def _checkout(self, *args, **kwargs) -> Optional[bool]:
         from dvc_data.hashfile.checkout import CheckoutError as _CheckoutError
         from dvc_data.hashfile.checkout import LinkError, PromptError
 
@@ -679,9 +679,9 @@ class Output:
         except LinkError as exc:
             raise CacheLinkError([exc.path])  # noqa: B904
         except _CheckoutError as exc:
-            raise CheckoutError(exc.paths)  # noqa: B904
+            raise CheckoutError(exc.paths, {})  # noqa: B904
 
-    def commit(self, filter_info=None):
+    def commit(self, filter_info=None) -> None:
         if not self.exists:
             raise self.DoesNotExistError(self)
 
@@ -721,7 +721,7 @@ class Output:
             )
             self.set_exec()
 
-    def _commit_granular_dir(self, filter_info):
+    def _commit_granular_dir(self, filter_info) -> Optional["HashFile"]:
         prefix = self.fs.path.parts(
             self.fs.path.relpath(filter_info, self.fs_path)
         )
@@ -850,14 +850,14 @@ class Output:
 
     def checkout(
         self,
-        force=False,
+        force: bool = False,
         progress_callback: "Callback" = DEFAULT_CALLBACK,
-        relink=False,
-        filter_info=None,
-        allow_missing=False,
-        checkpoint_reset=False,
+        relink: bool = False,
+        filter_info: Optional[str] = None,
+        allow_missing: bool = False,
+        checkpoint_reset: bool = False,
         **kwargs,
-    ):
+    ) -> Optional[Tuple[bool, Optional[bool]]]:
         if not self.use_cache:
             if progress_callback != DEFAULT_CALLBACK:
                 progress_callback.relative_update(

dvc/parsing/__init__.py

@@ -41,6 +41,8 @@ if TYPE_CHECKING:
 
     from dvc.repo import Repo
 
+from dvc.types import DictStrAny
+
 logger = logging.getLogger(__name__)
 
 STAGES_KWD = "stages"
@@ -53,7 +55,6 @@ DO_KWD = "do"
 DEFAULT_PARAMS_FILE = "params.yaml"
 
 JOIN = "@"
-DictStr = Dict[str, Any]
 
 
 class ResolveError(DvcException):
@@ -91,7 +92,7 @@ def _reraise_err(
 
 
 def check_syntax_errors(
-    definition: DictStr, name: str, path: str, where: str = "stages"
+    definition: DictStrAny, name: str, path: str, where: str = "stages"
 ):
     for key, d in definition.items():
         try:
@@ -110,8 +111,8 @@ def split_foreach_name(name: str) -> Tuple[str, Optional[str]]:
     return group, first(keys)
 
 
-def check_interpolations(data: DictStr, where: str, path: str):
-    def func(s: DictStr) -> None:
+def check_interpolations(data: DictStrAny, where: str, path: str):
+    def func(s: DictStrAny) -> None:
         if is_interpolated_string(s):
             raise ResolveError(
                 _format_preamble(f"'{where}'", path)
@@ -125,7 +126,7 @@ Definition = Union["ForeachDefinition", "EntryDefinition"]
 
 
 def make_definition(
-    resolver: "DataResolver", name: str, definition: DictStr, **kwargs
+    resolver: "DataResolver", name: str, definition: DictStrAny, **kwargs
 ) -> Definition:
     args = resolver, resolver.context, name, definition
     if FOREACH_KWD in definition:
@@ -224,7 +225,7 @@ class EntryDefinition:
         resolver: DataResolver,
         context: Context,
         name: str,
-        definition: DictStr,
+        definition: DictStrAny,
         where: str = STAGES_KWD,
     ):
         self.resolver = resolver
@@ -253,7 +254,7 @@ class EntryDefinition:
         except ContextError as exc:
             format_and_raise(exc, f"stage '{self.name}'", self.relpath)
 
-    def resolve_stage(self, skip_checks: bool = False) -> DictStr:
+    def resolve_stage(self, skip_checks: bool = False) -> DictStrAny:
         context = self.context
         name = self.name
         if not skip_checks:
@@ -301,7 +302,7 @@ class EntryDefinition:
 
     def _resolve(
         self, context: "Context", value: Any, key: str, skip_checks: bool
-    ) -> DictStr:
+    ) -> DictStrAny:
         try:
             return context.resolve(
                 value, skip_interpolation_checks=skip_checks, key=key
@@ -323,7 +324,7 @@ class ForeachDefinition:
         resolver: DataResolver,
         context: Context,
         name: str,
-        definition: DictStr,
+        definition: DictStrAny,
         where: str = STAGES_KWD,
     ):
         self.resolver = resolver
@@ -416,13 +417,13 @@ class ForeachDefinition:
     def _generate_name(self, key: str) -> str:
         return f"{self.name}{JOIN}{key}"
 
-    def resolve_all(self) -> DictStr:
+    def resolve_all(self) -> DictStrAny:
         return join(map(self.resolve_one, self.normalized_iterable))
 
-    def resolve_one(self, key: str) -> DictStr:
+    def resolve_one(self, key: str) -> DictStrAny:
         return self._each_iter(key)
 
-    def _each_iter(self, key: str) -> DictStr:
+    def _each_iter(self, key: str) -> DictStrAny:
         err_message = f"Could not find '{key}' in foreach group '{self.name}'"
         with reraise(KeyError, EntryNotFound(err_message)):
             value = self.normalized_iterable[key]

dvc/proc/process.py

@@ -5,7 +5,7 @@ import shlex
 import subprocess
 from contextlib import AbstractContextManager
 from dataclasses import asdict, dataclass
-from typing import List, Optional, TextIO, Union
+from typing import BinaryIO, List, Optional, Union
 
 from shortuuid import uuid
 
@@ -59,8 +59,8 @@ class ManagedProcess(AbstractContextManager):
         self.wdir = wdir
         self.name = name or uuid()
         self.returncode: Optional[int] = None
-        self._stdout: Optional[TextIO] = None
-        self._stderr: Optional[TextIO] = None
+        self._stdout: Optional[BinaryIO] = None
+        self._stderr: Optional[BinaryIO] = None
 
         self._run()
 
@@ -111,7 +111,7 @@ class ManagedProcess(AbstractContextManager):
         with open(self.pidfile_path, "w", encoding="utf-8") as fobj:
             fobj.write(str(self.pid))
 
-    def _run(self):
+    def _run(self) -> None:
         self._make_wdir()
         logger.debug(
             "Appending output to '%s'",

dvc/prompt.py

@@ -3,11 +3,12 @@
 import logging
 import sys
 from getpass import getpass
+from typing import Collection, Optional
 
 logger = logging.getLogger(__name__)
 
 
-def ask(prompt, limited_to=None):
+def ask(prompt: str, limited_to: Optional[Collection[str]] = None):
     if not sys.stdout.isatty():
         return None
 
@@ -29,7 +30,7 @@ def ask(prompt, limited_to=None):
         )
 
 
-def confirm(statement):
+def confirm(statement: str) -> bool:
     """Ask the user for confirmation about the specified statement.
 
     Args:
@@ -43,7 +44,7 @@ def confirm(statement):
     return answer and answer.startswith("y")
 
 
-def password(statement):
+def password(statement: str) -> str:
     """Ask the user for a password.
 
     Args:

dvc/repo/__init__.py

@@ -3,7 +3,15 @@ import os
 from collections import defaultdict
 from contextlib import contextmanager
 from functools import wraps
-from typing import TYPE_CHECKING, Callable, Iterable, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
 
 from dvc.exceptions import FileMissingError
 from dvc.exceptions import IsADirectoryError as DvcIsADirectoryError
@@ -17,9 +25,11 @@ if TYPE_CHECKING:
     from dvc.fs import FileSystem
     from dvc.fs.data import DataFileSystem
     from dvc.fs.dvc import DVCFileSystem
+    from dvc.lock import LockBase
     from dvc.machine import MachineManager
     from dvc.scm import Base, Git, NoSCM
     from dvc.stage import Stage
+    from dvc.types import DictStrAny
 
     from .experiments import Experiments
     from .index import Index
@@ -31,7 +41,7 @@ logger = logging.getLogger(__name__)
 @contextmanager
 def lock_repo(repo: "Repo"):
     # pylint: disable=protected-access
-    depth = repo._lock_depth
+    depth: int = repo._lock_depth
     repo._lock_depth += 1
 
     try:
@@ -95,13 +105,13 @@ class Repo:
         root_dir: Optional[str] = None,
         fs: Optional["FileSystem"] = None,
         uninitialized: bool = False,
-        scm: Optional["Base"] = None,
-    ):
+        scm: Optional[Union["Git", "NoSCM"]] = None,
+    ) -> Tuple[str, Optional[str], Optional[str]]:
         from dvc.fs import localfs
         from dvc.scm import SCM, SCMError
 
-        dvc_dir = None
-        tmp_dir = None
+        dvc_dir: Optional[str] = None
+        tmp_dir: Optional[str] = None
         try:
             root_dir = self.find_root(root_dir, fs)
             fs = fs or localfs
@@ -123,7 +133,7 @@ class Repo:
         assert root_dir
         return root_dir, dvc_dir, tmp_dir
 
-    def _get_database_dir(self, db_name):
+    def _get_database_dir(self, db_name: str) -> Optional[str]:
         # NOTE: by default, store SQLite-based remote indexes and state's
         # `links` and `md5s` caches in the repository itself to avoid any
         # possible state corruption in 'shared cache dir' scenario, but allow
@@ -151,15 +161,15 @@ class Repo:
 
     def __init__(
         self,
-        root_dir=None,
-        fs=None,
-        rev=None,
-        subrepos=False,
-        uninitialized=False,
-        config=None,
-        url=None,
-        repo_factory=None,
-        scm=None,
+        root_dir: Optional[str] = None,
+        fs: Optional["FileSystem"] = None,
+        rev: Optional[str] = None,
+        subrepos: bool = False,
+        uninitialized: bool = False,
+        config: Optional["DictStrAny"] = None,
+        url: Optional[str] = None,
+        repo_factory: Optional[Callable] = None,
+        scm: Optional["Base"] = None,
     ):
         from dvc.config import Config
         from dvc.data_cloud import DataCloud
@@ -184,6 +194,9 @@ class Repo:
             root_dir = "/"
             self._fs = GitFileSystem(scm=self._scm, rev=rev)
 
+        self.root_dir: str
+        self.dvc_dir: Optional[str]
+        self.tmp_dir: Optional[str]
         self.root_dir, self.dvc_dir, self.tmp_dir = self._get_repo_dirs(
             root_dir=root_dir,
             fs=self.fs,
@@ -191,15 +204,17 @@ class Repo:
             scm=scm,
         )
 
-        self.config = Config(self.dvc_dir, fs=self.fs, config=config)
+        self.config: Config = Config(self.dvc_dir, fs=self.fs, config=config)
         self._uninitialized = uninitialized
 
         # used by DVCFileSystem to determine if it should traverse subrepos
         self.subrepos = subrepos
 
         self.cloud = DataCloud(self)
-        self.stage = StageLoad(self)
+        self.stage: "StageLoad" = StageLoad(self)
 
+        self.lock: "LockBase"
+        self.odb: ODBManager
         if isinstance(self.fs, GitFileSystem) or not self.dvc_dir:
             self.lock = LockNoop()
             self.state = StateNoop()
@@ -230,13 +245,13 @@ class Repo:
             self._ignore()
 
         self.metrics = Metrics(self)
-        self.plots = Plots(self)
+        self.plots: Plots = Plots(self)
         self.params = Params(self)
 
         self.stage_collection_error_handler: Optional[
             Callable[[str, Exception], None]
         ] = None
-        self._lock_depth = 0
+        self._lock_depth: int = 0
 
     def __str__(self):
         return self.url or self.root_dir

dvc/repo/diff.py

@@ -1,17 +1,38 @@
 import logging
 import os
 from collections import defaultdict
-from typing import Dict, List
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Tuple,
+)
 
 from dvc.exceptions import PathMissingError
 from dvc.repo import locked
 from dvc.utils.collections import ensure_list
 
+if TYPE_CHECKING:
+    from dvc.fs import DVCFileSystem, FileSystem
+    from dvc.output import Output
+    from dvc.repo import Repo
+    from dvc_objects.obj import Object
+
 logger = logging.getLogger(__name__)
 
 
 @locked
-def diff(self, a_rev="HEAD", b_rev=None, targets=None):
+def diff(
+    self,
+    a_rev: str = "HEAD",
+    b_rev: Optional[str] = None,
+    targets: Optional[List[str]] = None,
+):
     """
     By default, it compares the workspace with the last commit's fs.
 
@@ -26,8 +47,7 @@ def diff(self, a_rev="HEAD", b_rev=None, targets=None):
     from dvc.fs.dvc import DVCFileSystem
 
     dvcfs = DVCFileSystem(repo=self)
-    targets = ensure_list(targets)
-    targets = [dvcfs.from_os_path(target) for target in targets]
+    targets = [dvcfs.from_os_path(target) for target in ensure_list(targets)]
 
     b_rev = b_rev if b_rev else "workspace"
     results = {}
@@ -98,7 +118,9 @@ def diff(self, a_rev="HEAD", b_rev=None, targets=None):
     return ret if any(ret.values()) else {}
 
 
-def _paths_checksums(repo, targets):
+def _paths_checksums(
+    repo: "Repo", targets: Optional[List[str]]
+) -> Dict[str, Optional[str]]:
     """
     A dictionary of checksums addressed by relpaths collected from
     the current fs outputs.
@@ -113,18 +135,20 @@ def _paths_checksums(repo, targets):
     return dict(_output_paths(repo, targets))
 
 
-def _output_paths(repo, targets):  # noqa: C901
+def _output_paths(  # noqa: C901
+    repo: "Repo", targets: Optional[List[str]]
+) -> Iterable[Tuple[str, Optional[str]]]:
     from dvc.fs import LocalFileSystem
     from dvc_data.hashfile.build import build
 
     on_working_fs = isinstance(repo.fs, LocalFileSystem)
 
-    def _exists(output):
+    def _exists(output: "Output") -> bool:
         if on_working_fs:
             return output.exists
         return True
 
-    def _to_path(output):
+    def _to_path(output: "Output") -> str:
         relparts = output.fs.path.relparts(output.fs_path)
         base = os.path.join(*relparts)
         if output.is_dir_checksum:
@@ -132,45 +156,55 @@ def _output_paths(repo, targets):  # noqa: C901
         return base
 
     for output in repo.index.outs:
-        if _exists(output):
-            yield_output = targets is None or any(
-                output.fs.path.isin_or_eq(output.fs_path, target)
-                for target in targets
+        if not _exists(output):
+            continue
+
+        yield_output = targets is None or any(
+            output.fs.path.isin_or_eq(output.fs_path, target)
+            for target in targets
+        )
+
+        if on_working_fs:
+            _, _, obj = build(
+                repo.odb.repo,
+                output.fs_path,
+                repo.odb.repo.fs,
+                "md5",
+                dry_run=True,
+                ignore=output.dvcignore,
             )
+            hash_info = obj.hash_info
+        else:
+            hash_info = output.hash_info
+            obj = output.get_obj()
 
-            if on_working_fs:
-                _, _, obj = build(
-                    repo.odb.repo,
-                    output.fs_path,
-                    repo.odb.repo.fs,
-                    "md5",
-                    dry_run=True,
-                    ignore=output.dvcignore,
-                )
-                hash_info = obj.hash_info
-            else:
-                hash_info = output.hash_info
-                obj = output.get_obj()
+        if yield_output:
+            yield _to_path(output), hash_info.value
 
-            if yield_output:
-                yield _to_path(output), hash_info.value
+        if not obj:
+            continue
 
-            if not obj:
-                continue
-
-            if output.is_dir_checksum and (
-                yield_output
-                or any(
+        if output.is_dir_checksum and (
+            yield_output
+            or (
+                isinstance(targets, list)
+                and any(
                     output.fs.path.isin(target, output.fs_path)
                     for target in targets
                 )
-            ):
-                yield from _dir_output_paths(
-                    output.fs, output.fs_path, obj, targets
-                )
+            )
+        ):
+            yield from _dir_output_paths(
+                output.fs, output.fs_path, obj, targets
+            )
 
 
-def _dir_output_paths(fs, fs_path, obj, targets=None):
+def _dir_output_paths(
+    fs: "FileSystem",
+    fs_path: str,
+    obj: "Object",
+    targets: Optional[Iterable[str]] = None,
+) -> Iterator[Tuple[str, str]]:
     base = os.path.join(*fs.path.relparts(fs_path))
     for key, _, oid in obj:
         fname = fs.path.join(fs_path, *key)
@@ -181,7 +215,7 @@ def _dir_output_paths(fs, fs_path, obj, targets=None):
             yield os.path.join(base, *key), oid.value
 
 
-def _filter_missing(dvcfs, paths):
+def _filter_missing(dvcfs: "DVCFileSystem", paths: Set[str]) -> Iterator[str]:
     for path in paths:
         fs_path = dvcfs.from_os_path(path)
         try:
@@ -198,9 +232,14 @@ def _filter_missing(dvcfs, paths):
             pass
 
 
-def _targets_to_paths(dvcfs, targets):
-    paths = []
-    missing = []
+def _targets_to_paths(
+    dvcfs: "DVCFileSystem", targets: List[str]
+) -> Tuple[List[str], List[str]]:
+    paths: List[str] = []
+    missing: List[str] = []
+
+    if not targets:
+        return paths, missing
 
     for target in targets:
         if dvcfs.exists(target):
@@ -211,7 +250,12 @@ def _targets_to_paths(dvcfs, targets):
     return paths, missing
 
 
-def _calculate_renamed(new, old, added, deleted):
+def _calculate_renamed(
+    new: Dict[str, Any],
+    old: Dict[str, Any],
+    added: Iterable[str],
+    deleted: Iterable[str],
+) -> List[Dict[str, Dict[str, str]]]:
     old_inverted: Dict[str, List[str]] = defaultdict(list)
     # It is needed to be dict of lists to cover cases
     # when repo has paths with same hash

dvc/repo/experiments/__init__.py

@@ -410,7 +410,7 @@ class Experiments:
     def get_exact_name(self, revs: Iterable[str]) -> Dict[str, Optional[str]]:
         """Returns preferred name for the specified revision.
 
-        Prefers tags, branches (heads), experiments in that orer.
+        Prefers tags, branches (heads), experiments in that order.
         """
         result: Dict[str, Optional[str]] = {}
         exclude = f"{EXEC_NAMESPACE}/*"

dvc/repo/experiments/apply.py

@@ -70,6 +70,7 @@ def apply(repo: "Repo", rev: str, force: bool = True, **kwargs):  # noqa: C901
     repo.scm.reset()
 
     if is_stash:
+        assert repo.tmp_dir is not None
         args_path = os.path.join(repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE)
         if os.path.exists(args_path):
             remove(args_path)

dvc/repo/experiments/executor/local.py

@@ -167,6 +167,7 @@ class TempDirExecutor(BaseLocalExecutor):
         wdir: Optional[str] = None,
         **kwargs,
     ):
+        assert repo.tmp_dir
         parent_dir: str = wdir or os.path.join(repo.tmp_dir, EXEC_TMP_DIR)
         os.makedirs(parent_dir, exist_ok=True)
         tmp_dir = mkdtemp(dir=parent_dir)

dvc/repo/experiments/init.py

@@ -283,6 +283,7 @@ def init(
             plots_key: compact(plots),
         },
     )
+    assert isinstance(stage, PipelineStage)
 
     with _disable_logging(), repo.scm_context(autostage=True, quiet=True):
         stage.dump(update_lock=False)
@@ -292,5 +293,4 @@ def init(
         if params:
             repo.scm_context.track_file(params)
 
-    assert isinstance(stage, PipelineStage)
     return stage, initialized_deps, initialized_out_dirs

dvc/repo/experiments/queue/base.py

@@ -112,6 +112,7 @@ class BaseStashQueue(ABC):
             failed_ref: Failed run Git stash ref for this queue.
         """
         self.repo = repo
+        assert self.repo.tmp_dir
         self.ref = ref
         self.failed_ref = failed_ref
 
@@ -129,10 +130,12 @@ class BaseStashQueue(ABC):
 
     @cached_property
     def pid_dir(self) -> str:
+        assert self.repo.tmp_dir is not None
         return os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR)
 
     @cached_property
     def args_file(self) -> str:
+        assert self.repo.tmp_dir is not None
         return os.path.join(self.repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE)
 
     @abstractmethod

dvc/repo/experiments/queue/celery.py

@@ -68,6 +68,7 @@ class LocalCeleryQueue(BaseStashQueue):
 
     @cached_property
     def wdir(self) -> str:
+        assert self.repo.tmp_dir is not None
         return os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, self.CELERY_DIR)
 
     @cached_property
@@ -321,7 +322,9 @@ class LocalCeleryQueue(BaseStashQueue):
                 fail_to_kill_entries[queue_entry] = rev
         return fail_to_kill_entries
 
-    def _mark_inactive_tasks_failure(self, remained_entries):
+    def _mark_inactive_tasks_failure(
+        self, remained_entries: Dict[QueueEntry, str]
+    ) -> None:
         remained_revs: List[str] = []
         running_ids = self._get_running_task_ids()
         logger.debug(f"Current running tasks ids: {running_ids}.")
@@ -343,7 +346,9 @@ class LocalCeleryQueue(BaseStashQueue):
         if remained_revs:
             raise CannotKillTasksError(remained_revs)
 
-    def _kill_entries(self, entries: Dict[QueueEntry, str], force: bool):
+    def _kill_entries(
+        self, entries: Dict[QueueEntry, str], force: bool
+    ) -> None:
         logger.debug(
             "Found active tasks: '%s' to kill",
             list(entries.values()),

dvc/repo/experiments/queue/tempdir.py

@@ -38,6 +38,7 @@ class TempDirQueue(WorkspaceQueue):
 
     @cached_property
     def _standalone_tmp_dir(self) -> str:
+        assert self.repo.tmp_dir is not None
         return os.path.join(self.repo.tmp_dir, _STANDALONE_TMP_DIR)
 
     @cached_property

dvc/repo/experiments/utils.py

@@ -51,6 +51,7 @@ def get_exp_rwlock(
     writes = writes or []
 
     cmd = " ".join(sys.argv)
+    assert repo.tmp_dir is not None
     path = os.path.join(repo.tmp_dir, EXEC_TMP_DIR)
     repo.fs.makedirs(path, exist_ok=True)
 

dvc/repo/metrics/show.py

@@ -7,7 +7,7 @@ from scmrepo.exceptions import SCMError
 from dvc.fs.dvc import DVCFileSystem
 from dvc.output import Output
 from dvc.repo import locked
-from dvc.repo.collect import StrPaths, collect
+from dvc.repo.collect import collect
 from dvc.scm import NoSCMError
 from dvc.utils import (
     as_posix,
@@ -25,7 +25,7 @@ def _is_metric(out: Output) -> bool:
     return bool(out.metric)
 
 
-def _to_fs_paths(metrics: List[Output]) -> StrPaths:
+def _to_fs_paths(metrics: List[Output]) -> List["str"]:
     result = []
     for out in metrics:
         if out.metric:

dvc/repo/plots/__init__.py

@@ -29,8 +29,10 @@ from dvc.utils.serialize import LOADERS
 from dvc.utils.threadpool import ThreadPoolExecutor
 
 if TYPE_CHECKING:
+    from dvc.fs import FileSystem
     from dvc.output import Output
     from dvc.repo import Repo
+    from dvc.types import DictAny, DictStrAny, StrPath
 
 dpath.options.ALLOW_EMPTY_STRING_KEYS = True
 
@@ -420,7 +422,12 @@ def _adjust_sources(fs, plot_props, config_dir):
 
 
 def _resolve_definitions(
-    fs, targets, props, config_path, definitions, onerror=None
+    fs: "FileSystem",
+    targets: List[str],
+    props: "DictAny",
+    config_path: "StrPath",
+    definitions: "DictStrAny",
+    onerror: Optional[Callable[[Any], Any]] = None,
 ):
     config_dir = fs.path.dirname(config_path)
     result: Dict[str, Dict] = {}

dvc/stage/__init__.py

@@ -4,7 +4,19 @@ import string
 from collections import defaultdict
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
 
 from funcy import project
 
@@ -37,6 +49,9 @@ from .utils import (
 
 if TYPE_CHECKING:
     from dvc.dvcfile import ProjectFile, SingleStageFile
+    from dvc.output import Output
+    from dvc.repo import Repo
+    from dvc.types import StrPath
     from dvc_data.hashfile.hash_info import HashInfo
     from dvc_objects.db import ObjectDB
 
@@ -44,9 +59,14 @@ logger = logging.getLogger(__name__)
 # Disallow all punctuation characters except hyphen and underscore
 INVALID_STAGENAME_CHARS = set(string.punctuation) - {"_", "-"}
 Env = Dict[str, str]
+ChangedEntries = Tuple[List[str], List[str], Optional[str]]
+
+_T = TypeVar("_T")
 
 
-def loads_from(cls, repo, path, wdir, data):
+def loads_from(
+    cls: Type[_T], repo: "Repo", path: str, wdir: str, data: Dict[str, Any]
+) -> _T:
     kw = {
         "repo": repo,
         "path": path,
@@ -74,7 +94,7 @@ class RawData:
     generated_from: Optional[str] = None
 
 
-def create_stage(cls, repo, path, external=False, **kwargs):
+def create_stage(cls: Type[_T], repo, path, external=False, **kwargs) -> _T:
     from dvc.dvcfile import check_dvcfile_path
 
     wdir = os.path.abspath(kwargs.get("wdir", None) or os.curdir)
@@ -97,7 +117,7 @@ def create_stage(cls, repo, path, external=False, **kwargs):
     return stage
 
 
-def restore_fields(stage):
+def restore_fields(stage: "Stage") -> None:
     from .exceptions import StageNotFound
 
     if not stage.dvcfile.exists():
@@ -227,12 +247,12 @@ class Stage(params.StageParams):
         return relpath(self.path)
 
     @property
-    def is_data_source(self):
+    def is_data_source(self) -> bool:
         """Whether the DVC file was created with `dvc add` or `dvc import`"""
         return self.cmd is None
 
     @property
-    def is_callback(self):
+    def is_callback(self) -> bool:
         """
         A callback stage is always considered as changed,
         so it runs on every `dvc repro` call.
@@ -240,7 +260,7 @@ class Stage(params.StageParams):
         return self.cmd and not any((self.deps, self.outs))
 
     @property
-    def is_import(self):
+    def is_import(self) -> bool:
         """Whether the DVC file was created with `dvc import`."""
         return not self.cmd and len(self.deps) == 1 and len(self.outs) == 1
 
@@ -253,7 +273,7 @@ class Stage(params.StageParams):
         return self.is_import and (not self.outs[0].hash_info)
 
     @property
-    def is_repo_import(self):
+    def is_repo_import(self) -> bool:
         if not self.is_import:
             return False
 
@@ -262,11 +282,11 @@ class Stage(params.StageParams):
         return isinstance(self.deps[0], RepoDependency)
 
     @property
-    def is_versioned_import(self):
+    def is_versioned_import(self) -> bool:
         return self.is_import and self.deps[0].fs.version_aware
 
     @property
-    def is_checkpoint(self):
+    def is_checkpoint(self) -> bool:
         """
         A stage containing checkpoint outs is always considered as changed
         since the checkpoint out is a circular dependency.
@@ -307,7 +327,7 @@ class Stage(params.StageParams):
             env.update(current)
         return env
 
-    def changed_deps(self):
+    def changed_deps(self) -> bool:
         if self.frozen:
             return False
 
@@ -317,7 +337,7 @@ class Stage(params.StageParams):
         return self._changed_deps()
 
     @rwlocked(read=["deps"])
-    def _changed_deps(self):
+    def _changed_deps(self) -> bool:
         for dep in self.deps:
             status = dep.status()
             if status:
@@ -331,7 +351,7 @@ class Stage(params.StageParams):
         return False
 
     @rwlocked(read=["outs"])
-    def changed_outs(self):
+    def changed_outs(self) -> bool:
         for out in self.outs:
             status = out.status()
             if status:
@@ -345,14 +365,14 @@ class Stage(params.StageParams):
 
         return False
 
-    def changed_stage(self):
+    def changed_stage(self) -> bool:
         changed = self.md5 != self.compute_md5()
         if changed:
             logger.debug(self._changed_stage_entry())
         return changed
 
     @rwlocked(read=["deps", "outs"])
-    def changed(self):
+    def changed(self) -> bool:
         is_changed = (
             # Short-circuit order: stage md5 is fast,
             # deps are expected to change
@@ -365,7 +385,7 @@ class Stage(params.StageParams):
         return is_changed
 
     @rwlocked(write=["outs"])
-    def remove_outs(self, ignore_remove=False, force=False):
+    def remove_outs(self, ignore_remove=False, force=False) -> None:
         """Used mainly for `dvc remove --outs` and :func:`Stage.reproduce`."""
         for out in self.outs:
             if (out.persist or out.checkpoint) and not force:
@@ -375,16 +395,16 @@ class Stage(params.StageParams):
             logger.debug("Removing output '%s' of %s.", out, self)
             out.remove(ignore_remove=ignore_remove)
 
-    def unprotect_outs(self):
+    def unprotect_outs(self) -> None:
         for out in self.outs:
             out.unprotect()
 
-    def ignore_remove_outs(self):
+    def ignore_remove_outs(self) -> None:
         for out in self.outs:
             out.ignore_remove()
 
     @rwlocked(write=["outs"])
-    def remove(self, force=False, remove_outs=True, purge=True):
+    def remove(self, force=False, remove_outs=True, purge=True) -> None:
         if remove_outs:
             self.remove_outs(ignore_remove=True, force=force)
         else:
@@ -408,7 +428,7 @@ class Stage(params.StageParams):
             out.ignore()
 
     @rwlocked(read=["deps"], write=["outs"])
-    def reproduce(self, interactive=False, **kwargs):
+    def reproduce(self, interactive=False, **kwargs) -> Optional["Stage"]:
         if not (kwargs.get("force", False) or self.changed()):
             if not isinstance(self, PipelineStage) and self.is_data_source:
                 logger.info("'%s' didn't change, skipping", self.addressing)
@@ -439,7 +459,7 @@ class Stage(params.StageParams):
         remote=None,
         no_download=None,
         jobs=None,
-    ):
+    ) -> None:
         if not (self.is_repo_import or self.is_import):
             raise StageUpdateError(self.relpath)
         update_import(
@@ -451,13 +471,13 @@ class Stage(params.StageParams):
             jobs=jobs,
         )
 
-    def reload(self):
+    def reload(self) -> "Stage":
         return self.dvcfile.stage
 
-    def dumpd(self, **kwargs):
+    def dumpd(self, **kwargs) -> Dict[str, Any]:
         return get_dump(self, **kwargs)
 
-    def compute_md5(self):
+    def compute_md5(self) -> Optional[str]:
         # `dvc add`ed files don't need stage md5
         if self.is_data_source and not (self.is_import or self.is_repo_import):
             m = None
@@ -527,18 +547,18 @@ class Stage(params.StageParams):
                 if old_out is not None:
                     out.merge_version_meta(old_out)
 
-    def ignore_outs(self):
+    def ignore_outs(self) -> None:
         for out in self.outs:
             out.ignore()
 
     @staticmethod
-    def _changed_entries(entries):
+    def _changed_entries(entries) -> List[str]:
         return [str(entry) for entry in entries if entry.workspace_status()]
 
-    def _changed_stage_entry(self):
+    def _changed_stage_entry(self) -> str:
         return f"'md5' of {self} changed."
 
-    def changed_entries(self):
+    def changed_entries(self) -> ChangedEntries:
         changed_deps = self._changed_entries(self.deps)
         changed_outs = self._changed_entries(self.outs)
         return (
@@ -548,7 +568,7 @@ class Stage(params.StageParams):
         )
 
     @rwlocked(write=["outs"])
-    def commit(self, allow_missing=False, filter_info=None):
+    def commit(self, allow_missing=False, filter_info=None) -> None:
         from dvc.output import OutputDoesNotExistError
 
         link_failures = []
@@ -572,7 +592,7 @@ class Stage(params.StageParams):
         allow_missing=False,
         no_download=False,
         **kwargs,
-    ):
+    ) -> None:
         if (self.cmd or self.is_import) and not self.frozen and not dry:
             self.remove_outs(ignore_remove=False, force=False)
 
@@ -610,28 +630,31 @@ class Stage(params.StageParams):
                 self.commit(allow_missing=allow_missing)
 
     @rwlocked(read=["deps"], write=["outs"])
-    def _run_stage(self, dry, force, **kwargs):
+    def _run_stage(self, dry, force, **kwargs) -> None:
         return run_stage(self, dry, force, **kwargs)
 
     @rwlocked(read=["deps"], write=["outs"])
-    def _sync_import(self, dry, force, jobs, no_download):
+    def _sync_import(self, dry, force, jobs, no_download) -> None:
         sync_import(self, dry, force, jobs, no_download)
 
     @rwlocked(read=["outs"])
-    def _check_missing_outputs(self):
+    def _check_missing_outputs(self) -> None:
         check_missing_outputs(self)
 
-    def filter_outs(self, fs_path):
+    def filter_outs(self, fs_path) -> Iterable["Output"]:
         def _func(o):
             return o.fs.path.isin_or_eq(fs_path, o.fs_path)
 
         return filter(_func, self.outs) if fs_path else self.outs
 
     @rwlocked(write=["outs"])
-    def checkout(self, allow_missing=False, **kwargs):
-        stats = defaultdict(list)
+    def checkout(
+        self, allow_missing: bool = False, **kwargs
+    ) -> Dict[str, List["StrPath"]]:
+        stats: Dict[str, List["StrPath"]] = defaultdict(list)
         if self.is_partial_import:
-            return {}
+            return stats
+
         for out in self.filter_outs(kwargs.get("filter_info")):
             key, outs = self._checkout(
                 out,
@@ -643,7 +666,7 @@ class Stage(params.StageParams):
         return stats
 
     @staticmethod
-    def _checkout(out, **kwargs):
+    def _checkout(out, **kwargs) -> Tuple[Optional[str], List[str]]:
         try:
             result = out.checkout(**kwargs)
             added, modified = result or (None, None)
@@ -654,8 +677,10 @@ class Stage(params.StageParams):
             return "failed", exc.target_infos
 
     @rwlocked(read=["deps", "outs"])
-    def status(self, check_updates=False, filter_info=None):
-        ret = []
+    def status(
+        self, check_updates: bool = False, filter_info: Optional[bool] = None
+    ) -> Dict[str, List[Union[str, Dict[str, str]]]]:
+        ret: List[Union[str, Dict[str, str]]] = []
         show_import = (
             self.is_repo_import or self.is_versioned_import
         ) and check_updates
@@ -668,7 +693,7 @@ class Stage(params.StageParams):
         return {self.addressing: ret} if ret else {}
 
     @staticmethod
-    def _status(entries):
+    def _status(entries: Iterable["Output"]) -> Dict[str, str]:
         ret = {}
 
         for entry in entries:
@@ -676,42 +701,42 @@ class Stage(params.StageParams):
 
         return ret
 
-    def _status_deps(self, ret):
+    def _status_deps(self, ret) -> None:
         deps_status = self._status(self.deps)
         if deps_status:
             ret.append({"changed deps": deps_status})
 
-    def _status_outs(self, ret, filter_info):
+    def _status_outs(self, ret, filter_info) -> None:
         filter_outs = self.filter_outs(filter_info)
         outs_status = self._status(filter_outs)
         if outs_status:
             ret.append({"changed outs": outs_status})
 
-    def _status_always_changed(self, ret):
+    def _status_always_changed(self, ret) -> None:
         if self.is_callback or self.always_changed or self.is_checkpoint:
             ret.append("always changed")
 
-    def _status_stage(self, ret):
+    def _status_stage(self, ret) -> None:
         if self.changed_stage():
             ret.append("changed checksum")
 
-    def already_cached(self):
+    def already_cached(self) -> bool:
         return (
             not self.changed_stage()
             and self.deps_cached()
             and self.outs_cached()
         )
 
-    def deps_cached(self):
+    def deps_cached(self) -> bool:
         return all(not dep.changed() for dep in self.deps)
 
-    def outs_cached(self):
+    def outs_cached(self) -> bool:
         return all(
             not out.changed_cache() if out.use_cache else not out.changed()
             for out in self.outs
         )
 
-    def get_all_files_number(self, filter_info=None):
+    def get_all_files_number(self, filter_info=None) -> int:
         return sum(
             out.get_files_number(filter_info)
             for out in self.filter_outs(filter_info)
@@ -731,7 +756,7 @@ class Stage(params.StageParams):
         return used_objs
 
     @staticmethod
-    def _check_can_merge(stage, ancestor_out=None):
+    def _check_can_merge(stage, ancestor_out=None) -> None:
         if isinstance(stage, PipelineStage):
             raise MergeError("unable to auto-merge pipeline stages")
 
@@ -746,7 +771,7 @@ class Stage(params.StageParams):
                 "unable to auto-merge DVC files with deleted outputs"
             )
 
-    def merge(self, ancestor, other, allowed=None):
+    def merge(self, ancestor, other, allowed=None) -> None:
         assert other
 
         if not other.outs:
@@ -768,21 +793,21 @@ class Stage(params.StageParams):
 
         self.outs[0].merge(ancestor_out, other.outs[0], allowed=allowed)
 
-    def dump(self, **kwargs):
+    def dump(self, **kwargs) -> None:
         self.dvcfile.dump(self, **kwargs)
 
 
 class PipelineStage(Stage):
-    def __init__(self, *args, name=None, **kwargs):
+    def __init__(self, *args, name: Optional[str] = None, **kwargs):
         super().__init__(*args, **kwargs)
         self.name = name
         self.cmd_changed = False
-        self.tracked_vars = {}
+        self.tracked_vars: Dict[str, Dict[str, Dict[str, str]]] = {}
 
     def __eq__(self, other):
         return super().__eq__(other) and self.name == other.name
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash((self.path_in_repo, self.name))
 
     @property
@@ -793,20 +818,24 @@ class PipelineStage(Stage):
             return self.name
         return f"{super().addressing}:{self.name}"
 
-    def reload(self):
+    def reload(self) -> Stage:
+        from dvc.dvcfile import ProjectFile
+
+        assert isinstance(self.dvcfile, ProjectFile)
+
         self.dvcfile._reset()  # pylint: disable=protected-access
         return self.dvcfile.stages[self.name]
 
-    def _status_stage(self, ret):
+    def _status_stage(self, ret) -> None:
         if self.cmd_changed:
             ret.append("changed command")
 
-    def changed_stage(self):
+    def changed_stage(self) -> bool:
         if self.cmd_changed:
             logger.debug(self._changed_stage_entry())
         return self.cmd_changed
 
-    def _changed_stage_entry(self):
+    def _changed_stage_entry(self) -> str:
         return f"'cmd' of {self} has changed."
 
     def merge(self, ancestor, other, allowed=None):

dvc/stage/loader.py

@@ -2,7 +2,7 @@ import logging
 from collections.abc import Mapping
 from copy import deepcopy
 from itertools import chain
-from typing import Any, Dict
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from funcy import get_in, lcat, once, project
 
@@ -18,13 +18,16 @@ from .exceptions import StageNameUnspecified, StageNotFound
 from .params import StageParams
 from .utils import fill_stage_dependencies, resolve_paths
 
+if TYPE_CHECKING:
+    from dvc.dvcfile import ProjectFile, SingleStageFile
+
 logger = logging.getLogger(__name__)
 
 
 class StageLoader(Mapping):
     def __init__(
         self,
-        dvcfile,
+        dvcfile: "ProjectFile",
         data,
         lockfile_data=None,
     ):
@@ -79,7 +82,9 @@ class StageLoader(Mapping):
             item.files = get_in(checksums, [key, path, item.PARAM_FILES])
 
     @classmethod
-    def load_stage(cls, dvcfile, name, stage_data, lock_data=None):
+    def load_stage(
+        cls, dvcfile: "ProjectFile", name, stage_data, lock_data=None
+    ):
         assert all([name, dvcfile, dvcfile.repo, dvcfile.path])
         assert stage_data and isinstance(stage_data, dict)
 
@@ -170,7 +175,12 @@ class StageLoader(Mapping):
 
 
 class SingleStageLoader(Mapping):
-    def __init__(self, dvcfile, stage_data, stage_text=None):
+    def __init__(
+        self,
+        dvcfile: "SingleStageFile",
+        stage_data: Dict[Any, str],
+        stage_text: Optional[str] = None,
+    ):
         self.dvcfile = dvcfile
         self.stage_data = stage_data or {}
         self.stage_text = stage_text
@@ -189,7 +199,9 @@ class SingleStageLoader(Mapping):
         )
 
     @classmethod
-    def load_stage(cls, dvcfile, d, stage_text):
+    def load_stage(
+        cls, dvcfile: "SingleStageFile", d: Dict[str, Any], stage_text: str
+    ) -> Stage:
         path, wdir = resolve_paths(
             dvcfile.repo.fs, dvcfile.path, d.get(Stage.PARAM_WDIR)
         )

dvc/stage/utils.py

@@ -1,6 +1,6 @@
 import os
 import pathlib
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 from funcy import concat, first, lsplit, rpartial
 
@@ -17,6 +17,7 @@ from .exceptions import (
 )
 
 if TYPE_CHECKING:
+    from dvc.dependency import Dependency, ParamsDependency
     from dvc.repo import Repo
 
     from . import PipelineStage, Stage
@@ -195,7 +196,7 @@ def resolve_paths(fs, path, wdir=None):
     return path, wdir
 
 
-def get_dump(stage, **kwargs):
+def get_dump(stage: "Stage", **kwargs):
     return {
         key: value
         for key, value in {
@@ -213,19 +214,21 @@ def get_dump(stage, **kwargs):
     }
 
 
-def split_params_deps(stage):
+def split_params_deps(
+    stage: "Stage",
+) -> Tuple[List["ParamsDependency"], List["Dependency"]]:
     from dvc.dependency import ParamsDependency
 
     return lsplit(rpartial(isinstance, ParamsDependency), stage.deps)
 
 
-def is_valid_name(name: str):
+def is_valid_name(name: str) -> bool:
     from . import INVALID_STAGENAME_CHARS
 
     return not INVALID_STAGENAME_CHARS & set(name)
 
 
-def prepare_file_path(kwargs):
+def prepare_file_path(kwargs) -> str:
     """Determine file path from the first output name.
 
     Used in creating .dvc files.
@@ -280,7 +283,7 @@ def check_stage_exists(
 
 def validate_kwargs(
     single_stage: bool = False, fname: Optional[str] = None, **kwargs
-):
+) -> Dict[str, Any]:
     """Prepare, validate and process kwargs passed from cli"""
     cmd = kwargs.get("cmd")
     if not cmd and not single_stage:

pyproject.toml

@@ -259,6 +259,7 @@ module = [
     "pygtrie.*",
     "pyinstrument",
     "pytest_docker.plugin",
+    "ruamel.*",
     "ruamel.yaml.*",
     "scmrepo.*",
     "tpi.*",