From e1254c9ece02247f409c686e699d08703aaa955e Mon Sep 17 00:00:00 2001 From: daniele <36171005+dtrifiro@users.noreply.github.com> Date: Sun, 15 Jan 2023 13:13:34 +0100 Subject: [PATCH] Add misc type hints (#8430) --- dvc/cli/command.py | 8 +- dvc/commands/machine.py | 2 +- dvc/commands/plots.py | 2 +- dvc/commands/queue/status.py | 10 +- dvc/config.py | 14 ++- dvc/dependency/repo.py | 38 +++++-- dvc/dvcfile.py | 7 +- dvc/exceptions.py | 4 +- dvc/external_repo.py | 12 +- dvc/fs/dvc.py | 29 +++-- dvc/ignore.py | 2 +- dvc/machine/__init__.py | 1 + dvc/output.py | 48 ++++---- dvc/parsing/__init__.py | 25 +++-- dvc/proc/process.py | 8 +- dvc/prompt.py | 7 +- dvc/repo/__init__.py | 55 +++++---- dvc/repo/diff.py | 130 ++++++++++++++-------- dvc/repo/experiments/__init__.py | 2 +- dvc/repo/experiments/apply.py | 1 + dvc/repo/experiments/executor/local.py | 1 + dvc/repo/experiments/init.py | 2 +- dvc/repo/experiments/queue/base.py | 3 + dvc/repo/experiments/queue/celery.py | 9 +- dvc/repo/experiments/queue/tempdir.py | 1 + dvc/repo/experiments/utils.py | 1 + dvc/repo/metrics/show.py | 4 +- dvc/repo/plots/__init__.py | 9 +- dvc/stage/__init__.py | 147 +++++++++++++++---------- dvc/stage/loader.py | 22 +++- dvc/stage/utils.py | 15 ++- pyproject.toml | 1 + 32 files changed, 399 insertions(+), 221 deletions(-) diff --git a/dvc/cli/command.py b/dvc/cli/command.py index 39a76984d..3d66c7449 100644 --- a/dvc/cli/command.py +++ b/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): diff --git a/dvc/commands/machine.py b/dvc/commands/machine.py index 0a4c112ad..c89f0ba30 100644 --- a/dvc/commands/machine.py +++ b/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 diff --git a/dvc/commands/plots.py b/dvc/commands/plots.py index dbf6e2ddd..9451db875 100644 --- a/dvc/commands/plots.py +++ b/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 diff --git a/dvc/commands/queue/status.py b/dvc/commands/queue/status.py index f5f7a9373..0ffcdea82 100644 --- a/dvc/commands/queue/status.py +++ b/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"], ] diff --git a/dvc/config.py b/dvc/config.py index f20fc59a3..2388c4d86 100644 --- a/dvc/config.py +++ b/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: diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py index 7cfd89259..7013c71ba 100644 --- a/dvc/dependency/repo.py +++ b/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 diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py index a1184ab5f..49eaf5100 100644 --- a/dvc/dvcfile.py +++ b/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 diff --git a/dvc/exceptions.py b/dvc/exceptions.py index ba5c4cbf3..dc9ac5df3 100644 --- a/dvc/exceptions.py +++ b/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 diff --git a/dvc/external_repo.py b/dvc/external_repo.py index c7a3eb4c4..b34032a93 100644 --- a/dvc/external_repo.py +++ b/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 diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index af618962f..165ce0984 100644 --- a/dvc/fs/dvc.py +++ b/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) diff --git a/dvc/ignore.py b/dvc/ignore.py index f89ec56a7..3c6108acf 100644 --- a/dvc/ignore.py +++ b/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__) diff --git a/dvc/machine/__init__.py b/dvc/machine/__init__.py index 1c2004d90..4fe54e854 100644 --- a/dvc/machine/__init__.py +++ b/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) diff --git a/dvc/output.py b/dvc/output.py index 2bfb98c33..d77b334fd 100644 --- a/dvc/output.py +++ b/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( diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py index 7e4c7a85f..056c72f3c 100644 --- a/dvc/parsing/__init__.py +++ b/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] diff --git a/dvc/proc/process.py b/dvc/proc/process.py index e66053acf..3d999ef5c 100644 --- a/dvc/proc/process.py +++ b/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'", diff --git a/dvc/prompt.py b/dvc/prompt.py index 73f054937..132a1b060 100644 --- a/dvc/prompt.py +++ b/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: diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 5ecb4ef57..69059dab4 100644 --- a/dvc/repo/__init__.py +++ b/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 diff --git a/dvc/repo/diff.py b/dvc/repo/diff.py index f5d330cc5..17aa2b309 100644 --- a/dvc/repo/diff.py +++ b/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 diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 52b2868e7..93c3c8b28 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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}/*" diff --git a/dvc/repo/experiments/apply.py b/dvc/repo/experiments/apply.py index 44df36e4c..feea6e620 100644 --- a/dvc/repo/experiments/apply.py +++ b/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) diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index 7072dbd25..e369aaca4 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/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) diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py index cd1d7cde1..4aceba92d 100644 --- a/dvc/repo/experiments/init.py +++ b/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 diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 49e21bd8e..5c14338d4 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/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 diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index e13246ec6..271b0aa1e 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/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()), diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index 4b469d6f2..4f534b88c 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/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 diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index 094d9b838..b1f7ed44a 100644 --- a/dvc/repo/experiments/utils.py +++ b/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) diff --git a/dvc/repo/metrics/show.py b/dvc/repo/metrics/show.py index c1b7fb73b..99c6dc7d7 100644 --- a/dvc/repo/metrics/show.py +++ b/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: diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py index 693015b8a..3e69f9afb 100644 --- a/dvc/repo/plots/__init__.py +++ b/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] = {} diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 4937eb67f..e3eb38fce 100644 --- a/dvc/stage/__init__.py +++ b/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): diff --git a/dvc/stage/loader.py b/dvc/stage/loader.py index ce9ee2e02..0ea7ce694 100644 --- a/dvc/stage/loader.py +++ b/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) ) diff --git a/dvc/stage/utils.py b/dvc/stage/utils.py index 7e3bf7adc..2be41340d 100644 --- a/dvc/stage/utils.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 5ebb91b78..873e862ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -259,6 +259,7 @@ module = [ "pygtrie.*", "pyinstrument", "pytest_docker.plugin", + "ruamel.*", "ruamel.yaml.*", "scmrepo.*", "tpi.*",