iterative/dvc

from dvc.cli import completion, formatter
from dvc.cli.command import CmdBase
from dvc.cli.utils import append_doc_link
from dvc.exceptions import DvcException
from dvc.log import logger

logger = logger.getChild(__name__)


class CmdFreezeBase(CmdBase):
    def _run(self, func, name):
        ret = 0
        for target in self.args.targets:
            try:
                func(target)
            except DvcException:
                logger.exception("failed to %s '%s'", name, target)
                ret = 1
        return ret


class CmdFreeze(CmdFreezeBase):
    def run(self):
        return self._run(self.repo.freeze, "freeze")


class CmdUnfreeze(CmdFreezeBase):
    def run(self):
        return self._run(self.repo.unfreeze, "unfreeze")


def add_parser(subparsers, parent_parser):
    FREEZE_HELP = "Freeze stages or .dvc files."
    freeze_parser = subparsers.add_parser(
        "freeze",
        parents=[parent_parser],
        description=append_doc_link(FREEZE_HELP, "freeze"),
        help=FREEZE_HELP,
        formatter_class=formatter.RawDescriptionHelpFormatter,
    )
    freeze_parser.add_argument(
        "targets", nargs="+", help="Stages or .dvc files to freeze"
    ).complete = completion.DVC_FILE
    freeze_parser.set_defaults(func=CmdFreeze)

    UNFREEZE_HELP = "Unfreeze stages or .dvc files."
    unfreeze_parser = subparsers.add_parser(
        "unfreeze",
        parents=[parent_parser],
        description=append_doc_link(UNFREEZE_HELP, "unfreeze"),
        help=UNFREEZE_HELP,
        formatter_class=formatter.RawDescriptionHelpFormatter,
    )
    unfreeze_parser.add_argument(
        "targets", nargs="+", help="Stages or .dvc files to unfreeze"
    ).complete = completion.DVC_FILE
    unfreeze_parser.set_defaults(func=CmdUnfreeze)