You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

checkpoints.py 2.2KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. """checkpoints logic."""
  2. import logging
  3. import os
  4. import shutil
  5. import sys
  6. import tempfile
  7. from config import Config as Conf
  8. from utils import setup_log, dl_file, unzip
  9. def main(_):
  10. """
  11. Start checkpoints main logic.
  12. :param _: None
  13. :return: None
  14. """
  15. if sum([1 for x in ["cm.lib", "mm.lib", "mn.lib"] if os.path.isfile(os.path.join(Conf.args['checkpoints'], x))]):
  16. Conf.log.info("Checkpoints Found In {}".format(Conf.args['checkpoints']))
  17. else:
  18. Conf.log.warn("Checkpoints Not Found In {}".format(Conf.args['checkpoints']))
  19. Conf.log.info("You Can Download Them Using : {} checkpoints download".format(sys.argv[0]))
  20. def download(_):
  21. """
  22. Start checkpoints download logic.
  23. :param _: None
  24. :return: None
  25. """
  26. Conf.log = setup_log(logging.DEBUG) if Conf.args['debug'] else setup_log()
  27. tempdir = tempfile.mkdtemp()
  28. cdn_url = Conf.checkpoints_cdn.format(Conf.checkpoints_version)
  29. temp_zip = os.path.join(tempdir, "{}.zip".format(Conf.checkpoints_version))
  30. try:
  31. Conf.log.info("Downloading {}".format(cdn_url))
  32. dl_file(Conf.checkpoints_cdn.format(Conf.checkpoints_version), temp_zip)
  33. if not os.path.exists(Conf.args['checkpoints']['checkpoints_path']):
  34. os.mkdir(Conf.args['checkpoints']['checkpoints_path'])
  35. Conf.log.info("Extracting {}".format(temp_zip))
  36. unzip(temp_zip, Conf.args['checkpoints']['checkpoints_path'])
  37. Conf.log.info("Moving Checkpoints To Final Location")
  38. for c in ("cm.lib", "mm.lib", "mn.lib"):
  39. if os.path.isfile(os.path.join(Conf.args['checkpoints']['checkpoints_path'], c)):
  40. os.remove(os.path.join(Conf.args['checkpoints']['checkpoints_path'], c))
  41. shutil.move(os.path.join(Conf.args['checkpoints']['checkpoints_path'], 'checkpoints', c), Conf.args['checkpoints']['checkpoints_path'])
  42. shutil.rmtree(os.path.join(Conf.args['checkpoints']['checkpoints_path'], 'checkpoints'))
  43. except Exception as e:
  44. Conf.log.error(e)
  45. Conf.log.error("Something Gone Bad Download Downloading The Checkpoints")
  46. shutil.rmtree(tempdir)
  47. sys.exit(1)
  48. shutil.rmtree(tempdir)
  49. Conf.log.info("Checkpoints Downloaded Successfully")