diff --git a/dvc/cli/command.py b/dvc/cli/command.py index 57c0dd318..39a76984d 100644 --- a/dvc/cli/command.py +++ b/dvc/cli/command.py @@ -13,7 +13,7 @@ class CmdBase(ABC): os.chdir(args.cd) - self.repo = Repo(uninitialized=self.UNINITIALIZED) + self.repo: "Repo" = Repo(uninitialized=self.UNINITIALIZED) self.config = self.repo.config self.args = args diff --git a/dvc/cli/parser.py b/dvc/cli/parser.py index dafa942f2..c626daf9d 100644 --- a/dvc/cli/parser.py +++ b/dvc/cli/parser.py @@ -14,6 +14,7 @@ from dvc.commands import ( config, daemon, dag, + data, data_sync, destroy, diff, @@ -88,6 +89,7 @@ COMMANDS = [ experiments, check_ignore, machine, + data, ] diff --git a/dvc/commands/data.py b/dvc/commands/data.py new file mode 100644 index 000000000..31fd45079 --- /dev/null +++ b/dvc/commands/data.py @@ -0,0 +1,174 @@ +import argparse +import logging +from typing import TYPE_CHECKING + +from funcy import compact, log_durations + +from dvc.cli.command import CmdBase +from dvc.cli.utils import append_doc_link, fix_subparsers +from dvc.ui import ui + +if TYPE_CHECKING: + from dvc.repo.data import Status as DataStatus + + +logger = logging.getLogger(__name__) + + +class CmdDataStatus(CmdBase): + COLORS = { + "not_in_cache": "red", + "committed": "green", + "uncommitted": "yellow", + "untracked": "cyan", + } + LABELS = { + "not_in_cache": "Not in cache", + "committed": "DVC committed changes", + "uncommitted": "DVC uncommitted changes", + "untracked": "Untracked files", + "unchanged": "DVC unchanged files", + } + HINTS = { + "not_in_cache": 'use "dvc pull ..." ' + "to update your local storage", + "committed": "git commit the corresponding dvc files " + "to update the repo", + "uncommitted": 'use "dvc commit ..." to track changes', + "untracked": 'use "git add ..." or ' + 'dvc add ..." to commit to git or to dvc', + "git_dirty": "there are {}changes not tracked by dvc, " + 'use "git status" to see', + } + + def _process_status(self, status: "DataStatus"): + """Flatten stage status, and filter empty stage status contents.""" + for stage, stage_status in status.items(): + items = stage_status + if isinstance(stage_status, dict): + items = { + file: state + for state, files in stage_status.items() + for file in files + } + if not items: + continue + yield stage, items + + def _show_status(self, status: "DataStatus") -> int: + git_info = status.pop("git") # type: ignore[misc] + result = dict(self._process_status(status)) + if not result: + no_changes = "No changes" + if git_info.get("is_empty", False): + no_changes += " in an empty git repo" + ui.write(f"{no_changes}.") + + for idx, (stage, stage_status) in enumerate(result.items()): + if idx: + ui.write() + + label = self.LABELS.get(stage, stage.capitalize() + " files") + header = f"{label}:" + color = self.COLORS.get(stage, "normal") + + ui.write(header) + if hint := self.HINTS.get(stage): + ui.write(f" ({hint})") + + if isinstance(stage_status, dict): + items = [ + ": ".join([state, file]) + for file, state in stage_status.items() + ] + else: + items = stage_status + + for item in items: + ui.write(f"\t[{color}]{item}[/]".expandtabs(8), styled=True) + + if (hint := self.HINTS.get("git_dirty")) and git_info.get("is_dirty"): + message = hint.format("other " if result else "") + ui.write(f"[blue]({message})[/]", styled=True) + return 0 + + def run(self) -> int: + with log_durations(logger.trace, "in data_status"): # type: ignore + status = self.repo.data_status( + granular=self.args.granular, + untracked_files=self.args.untracked_files, + with_dirs=self.args.with_dirs, + ) + + if not self.args.unchanged: + status.pop("unchanged") # type: ignore[misc] + if self.args.json: + status.pop("git") # type: ignore[misc] + ui.write_json(compact(status)) + return 0 + return self._show_status(status) + + +def add_parser(subparsers, parent_parser): + data_parser = subparsers.add_parser( + "data", + parents=[parent_parser], + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + data_subparsers = data_parser.add_subparsers( + dest="cmd", + help="Use `dvc data CMD --help` to display command-specific help.", + ) + fix_subparsers(data_subparsers) + + DATA_STATUS_HELP = ( + "Show changes between the last git commit, " + "the dvcfiles and the workspace." + ) + data_status_parser = data_subparsers.add_parser( + "status", + parents=[parent_parser], + description=append_doc_link(DATA_STATUS_HELP, "data/status"), + formatter_class=argparse.RawDescriptionHelpFormatter, + help=DATA_STATUS_HELP, + ) + data_status_parser.add_argument( + "--json", + action="store_true", + default=False, + help="Show output in JSON format.", + ) + data_status_parser.add_argument( + "--show-json", + action="store_true", + default=False, + dest="json", + help=argparse.SUPPRESS, + ) + data_status_parser.add_argument( + "--granular", + action="store_true", + default=False, + help="Show granular file-level info for DVC-tracked directories.", + ) + data_status_parser.add_argument( + "--unchanged", + action="store_true", + default=False, + help="Show unmodified DVC-tracked files.", + ) + data_status_parser.add_argument( + "--untracked-files", + choices=["no", "all"], + default="no", + const="all", + nargs="?", + help="Show untracked files.", + ) + data_status_parser.add_argument( + "--with-dirs", + action="store_true", + default=False, + help=argparse.SUPPRESS, + ) + data_status_parser.set_defaults(func=CmdDataStatus) diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index ce22ae95d..304d8f8c3 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -76,6 +76,8 @@ class Repo: from dvc.repo.status import status # type: ignore[misc] from dvc.repo.update import update # type: ignore[misc] + from .data import status as data_status # type: ignore[misc] + ls = staticmethod(_ls) get = staticmethod(_get) get_url = staticmethod(_get_url) diff --git a/dvc/repo/data.py b/dvc/repo/data.py new file mode 100644 index 000000000..38f6a1bc3 --- /dev/null +++ b/dvc/repo/data.py @@ -0,0 +1,245 @@ +import os +from collections import defaultdict +from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypedDict, cast + +if TYPE_CHECKING: + from scmrepo.base import Base + + from dvc.output import Output + from dvc.repo import Repo + from dvc_data.hashfile.db import HashFileDB + from dvc_data.hashfile.obj import HashFile + + +def _in_cache(obj: Optional["HashFile"], cache: "HashFileDB") -> bool: + from dvc_objects.errors import ObjectFormatError + + if not obj: + return False + if not obj.hash_info.value: + return False + + try: + cache.check(obj.hash_info.value) + return True + except (FileNotFoundError, ObjectFormatError): + return False + + +def _shallow_diff( + root: str, + old_obj: Optional["HashFile"], + new_obj: Optional["HashFile"], + cache: "HashFileDB", +) -> Dict[str, List[str]]: + # TODO: add support for shallow diff in dvc-data + # TODO: we may want to recursively do in_cache check + d = {} + + from dvc_data.objects.tree import Tree + + if isinstance(new_obj, Tree): + root = os.path.sep.join([root, ""]) + + if not _in_cache(old_obj, cache): + d["not_in_cache"] = [root] + + if old_obj is None and new_obj is None: + return d + if old_obj is None: + return {"added": [root], **d} + if new_obj is None: + return {"deleted": [root], **d} + if old_obj.hash_info != new_obj.hash_info: + return {"modified": [root], **d} + return {"unchanged": [root], **d} + + +def _granular_diff( + root: str, + old_obj: Optional["HashFile"], + new_obj: Optional["HashFile"], + cache: "HashFileDB", + with_dirs: bool = False, +) -> Dict[str, List[str]]: + from dvc_data.diff import ROOT + from dvc_data.diff import diff as odiff + from dvc_data.objects.tree import Tree + + def path_join(root: str, *paths: str) -> str: + if not isinstance(new_obj, Tree): + return root + return os.path.sep.join([root, *paths]) + + diff_data = odiff(old_obj, new_obj, cache) + drop_root = not with_dirs and isinstance(new_obj, Tree) + + output: Dict[str, List[str]] = defaultdict(list) + for state in ("added", "deleted", "modified", "unchanged"): + items = getattr(diff_data, state) + output[state].extend( + path_join(root, *item.new.key) + for item in items + if not (drop_root and item.new.key == ROOT) + ) + # TODO: PERF: diff is checking not_in_cache for each even if we only + # need it for the index. + # BUG: not_in_cache file also shows up as modified in staged and + # unstaged. We currently don't know if it is really modified. + output["not_in_cache"].extend( + path_join(root, *item.new.key) + for item in items + if not item.old.in_cache + and not (drop_root and item.new.key == ROOT) + and state != "added" + ) + return output + + +def _diff( + root: str, + old_obj: Optional["HashFile"], + new_obj: Optional["HashFile"], + cache: "HashFileDB", + granular: bool = False, + with_dirs: bool = False, +) -> Dict[str, List[str]]: + if granular: + return _granular_diff( + root, old_obj, new_obj, cache, with_dirs=with_dirs + ) + return _shallow_diff(root, old_obj, new_obj, cache) + + +class GitInfo(TypedDict, total=False): + staged: Dict[str, List[str]] + unstaged: Dict[str, List[str]] + untracked: List[str] + is_empty: bool + is_dirty: bool + + +def _git_info(scm: "Base", untracked_files: str = "all") -> GitInfo: + from scmrepo.exceptions import SCMError + + from dvc.scm import NoSCM + + if isinstance(scm, NoSCM): + return {} + + try: + scm.get_rev() + except SCMError: + empty_repo = True + else: + empty_repo = False + + staged, unstaged, untracked = scm.status(untracked_files=untracked_files) + # NOTE: order is important here. + return GitInfo( + staged=staged, + unstaged=unstaged, + untracked=untracked, + is_empty=empty_repo, + is_dirty=any([staged, unstaged, untracked]), + ) + + +def _diff_index_to_wtree(repo: "Repo", **kwargs: Any) -> Dict[str, List[str]]: + from dvc_data.build import build + + unstaged_diff = defaultdict(list) + for out in repo.index.outs: + out = cast("Output", out) + if not out.use_cache: + continue + + try: + _, _, new = build( + out.odb, + out.fs_path, + out.fs, + out.fs.PARAM_CHECKSUM, + ignore=out.dvcignore, + dry_run=True, + ) + except FileNotFoundError: + new = None + + cache = repo.odb.local + root = str(out) + old = out.get_obj() + d = _diff(root, old, new, cache, **kwargs) + for state, items in d.items(): + if not items: + continue + unstaged_diff[state].extend(items) + return unstaged_diff + + +def _diff_head_to_index( + repo: "Repo", head: str = "HEAD", **kwargs: Any +) -> Dict[str, List[str]]: + # we need to store objects from index and the HEAD to diff later + objs: Dict[str, Dict[str, "HashFile"]] = defaultdict(dict) + staged_diff = defaultdict(list) + for rev in repo.brancher(revs=[head]): + for out in repo.index.outs: + out = cast("Output", out) + if not out.use_cache: + continue + + root = str(out) + typ = "index" if rev == "workspace" else head + objs[root][typ] = out.get_obj() + + cache = repo.odb.local + for root, obj_d in objs.items(): + old = obj_d.get(head, None) + new = obj_d.get("index", None) + d = _diff(root, old, new, cache, **kwargs) + for state, items in d.items(): + if not items: + continue + staged_diff[state].extend(items) + + return staged_diff + + +class Status(TypedDict): + not_in_cache: List[str] + committed: Dict[str, Any] + uncommitted: Dict[str, Any] + untracked: List[str] + unchanged: List[str] + git: GitInfo + + +def status(repo: "Repo", untracked_files: str = "no", **kwargs: Any) -> Status: + from scmrepo.exceptions import SCMError + + from dvc.scm import NoSCMError + + head = kwargs.pop("head", "HEAD") + uncommitted_diff = _diff_index_to_wtree(repo, **kwargs) + not_in_cache = uncommitted_diff.pop("not_in_cache", []) + unchanged = set(uncommitted_diff.pop("unchanged", [])) + + try: + committed_diff = _diff_head_to_index(repo, head=head, **kwargs) + except (SCMError, NoSCMError): + committed_diff = {} + else: + # we don't care about not-in-cache between the head and the index. + committed_diff.pop("not_in_cache", None) + unchanged &= set(committed_diff.pop("unchanged", [])) + + git_info = _git_info(repo.scm, untracked_files=untracked_files) + return Status( + not_in_cache=not_in_cache, + committed=committed_diff, + uncommitted=uncommitted_diff, + untracked=git_info.get("untracked", []), + unchanged=list(unchanged), + git=git_info, + ) diff --git a/tests/unit/command/test_compat_flag.py b/tests/unit/command/test_compat_flag.py index 996121bbf..f5aa78eb2 100644 --- a/tests/unit/command/test_compat_flag.py +++ b/tests/unit/command/test_compat_flag.py @@ -34,6 +34,7 @@ def _id_gen(val) -> str: (["plots", "diff", "--show-json"], "json"), (["exp", "list", "--names-only"], "name_only"), (["stage", "list", "--names-only"], "name_only"), + (["data", "status", "--json"], "json"), ], ids=_id_gen, )