dvcpatch.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. """
  2. Support code for our custom DVC remote.
  3. """
  4. import logging
  5. from urllib.parse import urlparse
  6. import hashlib
  7. from dvc.remote.base import BaseRemoteTree
  8. from dvc.output.base import BaseOutput
  9. from dvc.dependency.base import BaseDependency
  10. from . import tracking
  11. _log = logging.getLogger('dvc.bgpatch')
  12. class PGRemoteTree(BaseRemoteTree):
  13. """
  14. PG status remote
  15. """
  16. scheme = 'pgstat'
  17. PARAM_CHECKSUM = 'md5'
  18. def __init__(self, *args, **kwargs):
  19. super().__init__(*args, **kwargs)
  20. def get_file_hash(self, path_info):
  21. _log.debug('checksum from %s', path_info)
  22. status = tracking.stage_status(path_info.bucket)
  23. h = hashlib.md5()
  24. h.update(status.encode('utf-8'))
  25. return h.hexdigest()
  26. def copy(self, from_info, to_info):
  27. _log.debug('copy from %s', from_info)
  28. _log.debug('copy to %s', to_info)
  29. raise NotImplementedError()
  30. def exists(self, path_info):
  31. _log.debug('exists? %s', path_info)
  32. return tracking.stage_exists(path_info.bucket)
  33. def remove(self, path_info):
  34. _log.info('asked to remove %s, ignoring', path_info)
  35. def _download(self, from_info, to_info, name, no_progress_bar):
  36. _log.info('download requested for %s', from_info)
  37. raise NotImplementedError()
  38. class PGOutput(BaseOutput):
  39. TREE_CLS = PGRemoteTree
  40. class PGDep(BaseDependency, PGOutput):
  41. pass
  42. def patch():
  43. "Patch DVC to include our classes"
  44. import dvc.output, dvc.dependency, dvc.config
  45. dvc.output.OUTS.append(PGOutput)
  46. dvc.output.OUTS_MAP['pgstat'] = PGOutput
  47. dvc.dependency.DEPS.append(PGDep)
  48. dvc.dependency.DEP_MAP['pgstat'] = PGDep
  49. # from dvc.cache import Cache, _make_remote_property
  50. # Cache.pgstat = _make_remote_property('pgstat')
  51. # dvc.config.Config.SCHEMA['cache']['pgstat'] = str
  52. # dvc.config.Config.COMPILED_SCHEMA = dvc.config.Schema(dvc.config.Config.SCHEMA)