Tooling for managing asset compression, storage, and retrieval
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.

832 lines
28 KiB

2 years ago
2 years ago
  1. #!/usr/bin/env python3
  2. import asyncio
  3. import collections.abc
  4. import hashlib
  5. import io
  6. import json
  7. import logging
  8. import os
  9. import pathlib
  10. import platform
  11. import sys
  12. import tempfile
  13. from typing import List, Dict, Callable
  14. import click
  15. from minio import Minio, InvalidResponseError
  16. from minio.error import S3Error
  17. from acm.config import VERSION, get_default_config
  18. from acm.logging import setup_basic_logging, update_logging_level
  19. from acm.utility import get_string_sha256sum
  20. # Size of the buffer to read files with
  21. BUF_SIZE = 4096
  22. LOG = setup_basic_logging("acm")
  23. ###########
  24. # AsyncIO #
  25. ###########
  26. async def run_command_shell(
  27. command,
  28. stdout=asyncio.subprocess.PIPE,
  29. stderr=asyncio.subprocess.PIPE,
  30. on_success: List[Callable] = [()]
  31. ):
  32. """Run command in subprocess (shell).
  33. Note:
  34. This can be used if you wish to execute e.g. "copy"
  35. on Windows, which can only be executed in the shell.
  36. """
  37. process = await asyncio.create_subprocess_shell(
  38. command, stdout=stdout, stderr=stderr
  39. )
  40. process_stdout, process_stderr = await process.communicate()
  41. if process.returncode == 0:
  42. for success_callable in on_success:
  43. success_callable()
  44. if stdout != asyncio.subprocess.DEVNULL:
  45. result = process_stdout.decode().strip()
  46. return result
  47. else:
  48. return None
  49. def make_chunks(tasks, chunk_size):
  50. """Yield successive chunk_size-sized chunks from tasks.
  51. Note:
  52. Taken from https://stackoverflow.com/a/312464
  53. modified for python 3 only
  54. """
  55. for i in range(0, len(tasks), chunk_size):
  56. yield tasks[i: i + chunk_size]
  57. def run_asyncio_commands(tasks, max_concurrent_tasks=0):
  58. """Run tasks asynchronously using asyncio and return results.
  59. If max_concurrent_tasks are set to 0, no limit is applied.
  60. Note:
  61. By default, Windows uses SelectorEventLoop, which does not support
  62. subprocesses. Therefore ProactorEventLoop is used on Windows.
  63. https://docs.python.org/3/library/asyncio-eventloops.html#windows
  64. """
  65. all_results = []
  66. if max_concurrent_tasks == 0:
  67. chunks = [tasks]
  68. num_chunks = len(chunks)
  69. else:
  70. chunks = make_chunks(tasks=tasks, chunk_size=max_concurrent_tasks)
  71. num_chunks = len(
  72. list(make_chunks(tasks=tasks, chunk_size=max_concurrent_tasks)))
  73. if asyncio.get_event_loop().is_closed():
  74. asyncio.set_event_loop(asyncio.new_event_loop())
  75. if platform.system() == "Windows":
  76. asyncio.set_event_loop(asyncio.ProactorEventLoop())
  77. loop = asyncio.get_event_loop()
  78. chunk = 1
  79. for tasks_in_chunk in chunks:
  80. commands = asyncio.gather(*tasks_in_chunk)
  81. results = loop.run_until_complete(commands)
  82. all_results += results
  83. chunk += 1
  84. loop.close()
  85. return all_results
  86. ###########
  87. # Helpers #
  88. ###########
  89. def update(d, u):
  90. for k, v in u.items():
  91. if isinstance(v, collections.abc.Mapping):
  92. d[k] = update(d.get(k, {}), v)
  93. else:
  94. d[k] = v
  95. return d
  96. def get_metadata_name(key):
  97. return METADATA_PREFIX + 'SHA256SUM'.capitalize()
  98. def get_clean_stdin_iterator(stdin_stream):
  99. return (line for line in [line.strip() for line in stdin_stream if line.strip() != ''])
  100. def strip_prefix(prefix: str, file: str) -> str:
  101. if file.startswith(prefix):
  102. return file.replace(prefix, '')
  103. return file
  104. def get_file_identity(ctx_obj, file):
  105. if 'REMOVE_PREFIX' in ctx_obj and ctx_obj['REMOVE_PREFIX'] is not None:
  106. path = strip_prefix(ctx_obj['REMOVE_PREFIX'], file)
  107. else:
  108. path = file
  109. if os.pathsep != '/':
  110. path = '/'.join(path.split(os.pathsep))
  111. return path
  112. def list_s3_dir(s3: Minio, bucket: str, prefix: str) -> List[str]:
  113. found_files = []
  114. for obj in s3.list_objects_v2(bucket, prefix=prefix):
  115. if obj.is_dir:
  116. found_files.extend(list_s3_dir(s3, bucket, obj.object_name))
  117. else:
  118. found_files.append(obj.object_name)
  119. return found_files
  120. def get_s3_client(config: any) -> Minio:
  121. host = config['host']
  122. secure = config['secure']
  123. access_key = config['access']
  124. secret_key = config['secret']
  125. return Minio(host, secure=secure, access_key=access_key, secret_key=secret_key)
  126. def prep_s3(ctx):
  127. s3_config = ctx.obj['CONFIG']['s3']
  128. s3_bucket = ctx.obj['CONTEXT']
  129. s3 = get_s3_client(s3_config)
  130. if not s3.bucket_exists(s3_bucket):
  131. s3.make_bucket(s3_bucket)
  132. return s3_bucket, s3
  133. def get_file_sha256sum(stored_data, profile, file):
  134. stored_file_hash = stored_data['sha256sum']
  135. stored_profile_hash = stored_data['profileHash']
  136. sha256sum = hashlib.sha256()
  137. with open(file, 'rb') as f:
  138. for byte_block in iter(lambda: f.read(BUF_SIZE), b""):
  139. sha256sum.update(byte_block)
  140. calculated_file_hash = sha256sum.hexdigest()
  141. return stored_profile_hash, stored_file_hash, calculated_file_hash
  142. def add_nested_key(config: Dict[str, any], path: List[str], value: str) -> bool:
  143. target = path[0].lower()
  144. if len(path) == 1:
  145. config[target] = value
  146. return True
  147. else:
  148. if target not in config:
  149. config[target] = {}
  150. add_nested_key(config[target], path[1:], value)
  151. return False
  152. def read_env_config(prefix, separator='__') -> any:
  153. prefix = prefix+separator
  154. env_config = {}
  155. environment_variables = [
  156. env for env in os.environ.keys() if env.startswith(prefix)]
  157. for env in environment_variables:
  158. path = env[len(prefix):].split('__')
  159. add_nested_key(env_config, path, os.environ[env])
  160. return env_config
  161. def load_config(path: str) -> any:
  162. combined_config = {}
  163. with open(
  164. os.path.join(
  165. os.path.dirname(os.path.realpath(__file__)),
  166. 'acm-config-default.json'),
  167. 'r') as combined_config_file:
  168. combined_config = json.load(combined_config_file)
  169. config = {}
  170. with open(path, 'r') as config_file:
  171. config = json.load(config_file)
  172. # Setup concurrency
  173. if 'concurrency' in config:
  174. config['concurrency'] = abs(int(config['concurrency']))
  175. else:
  176. config['concurrency'] = 0
  177. update(combined_config, config)
  178. update(combined_config, read_env_config('ACM'))
  179. # Calculate profiles hash
  180. profile_hashes = {}
  181. profile_hashes['all'] = get_string_sha256sum(
  182. json.dumps(combined_config['profiles']))
  183. for profile in combined_config['profiles'].keys():
  184. profile_hashes[profile] = get_string_sha256sum(
  185. json.dumps(combined_config['profiles'][profile]))
  186. combined_config['profileHashes'] = profile_hashes
  187. return combined_config
  188. @click.group()
  189. @click.option('-d', '--debug/--no-debug', default=False)
  190. @click.option('-c', '--config', default=lambda: os.path.join(os.getcwd(), 'acm-config.json'), show_default=True)
  191. @click.option('-s', '--stdin/--no-stdin', default=False)
  192. @click.option('--remove-prefix', default=None)
  193. @click.option('--add-prefix', default=None)
  194. @click.pass_context
  195. def cli(ctx, debug, config, stdin, remove_prefix, add_prefix):
  196. ctx.ensure_object(dict)
  197. # Propagate the global configs
  198. ctx.obj['DEBUG'] = debug
  199. # ctx.obj['CONFIG'] = load_config(config)
  200. ctx.obj['READ_STDIN'] = stdin
  201. ctx.obj['REMOVE_PREFIX'] = remove_prefix
  202. ctx.obj['ADD_PREFIX'] = add_prefix
  203. if debug:
  204. update_logging_level(3, LOG)
  205. # Reduce the logging noise for library loggers
  206. update_logging_level(0, "asyncio")
  207. ####################
  208. # Generic Commands #
  209. ####################
  210. @cli.command(name="config")
  211. @click.pass_context
  212. def print_config(ctx):
  213. """
  214. Print the configuration
  215. """
  216. print(json.dumps(ctx.obj['CONFIG'], indent=2, sort_keys=True))
  217. @cli.command(name="default-config")
  218. @click.argument('profile', default="all")
  219. @click.pass_context
  220. def print_default_config(ctx, profile):
  221. """
  222. Print the configuration
  223. """
  224. if profile == "all":
  225. print(get_default_config().json(exclude_none=True, indent=2, sort_keys=True))
  226. else:
  227. config = get_default_config()
  228. profile_names = config.get_profile_names()
  229. if profile in profile_names:
  230. print(config.get_profile(profile).json(exclude_none=True, indent=2, sort_keys=True))
  231. else:
  232. print(f"Profile \"{profile}\" is not in {profile_names}")
  233. ###############################
  234. # S3 Storage Focused Commands #
  235. ###############################
  236. @cli.command(name="list")
  237. @click.option('--sha256sum/--no-sha256sum', default=False)
  238. @click.option('--suffix', default=None)
  239. @click.option('-x', '--context', required=True)
  240. @click.option('--print-identity/--no-print-identity', default=False)
  241. @click.pass_context
  242. def list_files(ctx, context, sha256sum, suffix, print_identity):
  243. """
  244. List all file object in a bucket
  245. """
  246. ctx.obj['CONTEXT'] = context
  247. s3_config = ctx.obj['CONFIG']['s3']
  248. s3_bucket = ctx.obj['CONTEXT']
  249. LOG.debug(f"connecting to s3 {s3_config['']}")
  250. s3 = get_s3_client(s3_config)
  251. if not s3.bucket_exists(s3_bucket):
  252. s3.make_bucket(s3_bucket)
  253. found_files: List[str] = []
  254. found_objects: List[str] = []
  255. for obj in s3.list_objects_v2(s3_bucket, recursive=False):
  256. if obj.is_dir:
  257. found_objects.extend(list_s3_dir(s3, s3_bucket, obj.object_name))
  258. else:
  259. found_objects.append(obj.object_name)
  260. for obj in found_objects:
  261. file = obj
  262. if 'REMOVE_PREFIX' in ctx.obj and ctx.obj['REMOVE_PREFIX'] is not None:
  263. file = os.path.join(ctx.obj['REMOVE_PREFIX'], file)
  264. if suffix is not None and suffix in file:
  265. file = file.replace(suffix, '')
  266. file = file.strip()
  267. if sha256sum:
  268. file_object = s3.get_object(s3_bucket, obj)
  269. stored_data = json.load(file_object)
  270. sha256sum_value = stored_data['sha256sum']
  271. file = f'{sha256sum_value} {file}'
  272. elif print_identity:
  273. file_object = s3.get_object(s3_bucket, obj)
  274. stored_data = json.load(file_object)
  275. found_files.append(stored_data['storedAssetIdentity'])
  276. else:
  277. found_files.append(file)
  278. print(os.linesep.join(found_files))
  279. @cli.command(name="match")
  280. @click.option('-x', '--context', required=True)
  281. @click.option('--print-identity/--no-print-identity', default=False)
  282. @click.option('-p', '--profile', default='all')
  283. @click.argument('files', nargs=-1)
  284. @click.pass_context
  285. def check_matched_files_hashes(ctx, context, print_identity, profile, files):
  286. """
  287. List all files that have matching stored sha256sum and profile hash
  288. """
  289. ctx.obj['CONTEXT'] = context
  290. s3_bucket, s3 = prep_s3(ctx)
  291. matching_files: List[str] = []
  292. if ctx.obj['READ_STDIN']:
  293. files = get_clean_stdin_iterator(click.get_text_stream('stdin'))
  294. for file in files:
  295. file_identity = f'{get_file_identity(ctx.obj, file)}.json'
  296. try:
  297. file_object = s3.get_object(s3_bucket, file_identity)
  298. stored_data = json.load(file_object)
  299. stored_profile_hash, stored_file_hash, calculated_file_hash = get_file_sha256sum(
  300. stored_data, profile, file)
  301. if calculated_file_hash == stored_file_hash \
  302. and ctx.obj['CONFIG']['profileHashes'][profile] == stored_profile_hash:
  303. if print_identity:
  304. matching_files.append(stored_data['storedAssetIdentity'])
  305. else:
  306. matching_files.append(file)
  307. except S3Error as e:
  308. if e.code == "NoSuchKey":
  309. continue
  310. else:
  311. LOG.error(e)
  312. except ValueError or InvalidResponseError as e:
  313. LOG.error(f'ERROR: {file} {e}')
  314. print(os.linesep.join(matching_files))
  315. @cli.command(name="check")
  316. @click.option('-x', '--context', required=True)
  317. @click.option('-p', '--profile', default='all')
  318. @click.argument('files', nargs=-1)
  319. @click.pass_context
  320. def check_changed_files_hashes(ctx, context, profile, files):
  321. """
  322. List all files that do not have a matching sha256sum or profile hash
  323. """
  324. ctx.obj['CONTEXT'] = context
  325. s3_bucket, s3 = prep_s3(ctx)
  326. changed_files: List[str] = []
  327. if ctx.obj['READ_STDIN']:
  328. files = get_clean_stdin_iterator(click.get_text_stream('stdin'))
  329. for file in files:
  330. file_identity = f'{get_file_identity(ctx.obj, file)}.json'
  331. try:
  332. file_object = s3.get_object(s3_bucket, file_identity)
  333. stored_data = json.load(file_object)
  334. stored_profile_hash, stored_file_hash, calculated_file_hash = get_file_sha256sum(
  335. stored_data, profile, file)
  336. if calculated_file_hash != stored_file_hash \
  337. or ctx.obj['CONFIG']['profileHashes'][profile] != stored_profile_hash:
  338. changed_files.append(file)
  339. except S3Error as e:
  340. if e.code == "NoSuchKey":
  341. changed_files.append(file)
  342. else:
  343. LOG.error(e)
  344. except ValueError or InvalidResponseError as e:
  345. LOG.error(f'ERROR: {file} {e}')
  346. print(os.linesep.join(changed_files))
  347. @cli.command(name="update")
  348. @click.option('-x', '--context', required=True)
  349. @click.option('--input-and-identity/--no-input-and-identity', default=False)
  350. @click.option('-p', '--profile', default='all')
  351. @click.argument('files', nargs=-1)
  352. @click.pass_context
  353. def update_changed_files_hashes(ctx, context, input_and_identity, profile, files):
  354. """
  355. Store new data objects for the provided files
  356. """
  357. ctx.obj['CONTEXT'] = context
  358. s3_bucket, s3 = prep_s3(ctx)
  359. updated_files: List[str] = []
  360. if ctx.obj['READ_STDIN']:
  361. files = get_clean_stdin_iterator(click.get_text_stream('stdin'))
  362. for file in files:
  363. identity = None
  364. if input_and_identity:
  365. file, identity = file.split('\t')
  366. file_identity = f'{get_file_identity(ctx.obj, file)}.json'
  367. try:
  368. sha256sum = hashlib.sha256()
  369. with open(file, 'rb') as f:
  370. for byte_block in iter(lambda: f.read(BUF_SIZE), b''):
  371. sha256sum.update(byte_block)
  372. calculated_file_hash = sha256sum.hexdigest()
  373. object_data = {
  374. "sourcePath": file,
  375. "storedAssetIdentity": identity,
  376. "identity": file_identity,
  377. "sha256sum": calculated_file_hash,
  378. "profileHash": ctx.obj['CONFIG']['profileHashes'][profile]
  379. }
  380. with io.BytesIO(json.dumps(object_data, sort_keys=True, indent=None).encode('utf-8')) as data:
  381. data.seek(0, os.SEEK_END)
  382. data_length = data.tell()
  383. data.seek(0)
  384. s3.put_object(
  385. s3_bucket,
  386. file_identity,
  387. data,
  388. data_length,
  389. content_type="application/json",
  390. metadata={}
  391. )
  392. updated_files.append(file)
  393. except ValueError or InvalidResponseError as e:
  394. LOG.error(f'ERROR: {file} {e}')
  395. print(os.linesep.join(updated_files))
  396. @cli.command(name="store")
  397. @click.option('-x', '--context', required=True)
  398. @click.argument('files', nargs=-1)
  399. @click.pass_context
  400. def store_files(ctx, context, files):
  401. """
  402. Store specified files in a <context> bucket for retrieval.
  403. """
  404. ctx.obj['CONTEXT'] = context
  405. s3_bucket, s3 = prep_s3(ctx)
  406. stored_files: List[str] = []
  407. if ctx.obj['READ_STDIN']:
  408. files = get_clean_stdin_iterator(click.get_text_stream('stdin'))
  409. for file in files:
  410. file_identity = get_file_identity(ctx.obj, file)
  411. try:
  412. s3.fput_object(
  413. s3_bucket,
  414. file_identity,
  415. file,
  416. content_type="application/octet-stream"
  417. )
  418. if 'ADD_PREFIX' in ctx.obj and ctx.obj['ADD_PREFIX'] is not None:
  419. stored_files.append(os.path.join(
  420. ctx.obj['ADD_PREFIX'], file_identity))
  421. else:
  422. stored_files.append(file)
  423. except InvalidResponseError as e:
  424. LOG.error(f'ERROR: {file} {e}', file=sys.stderr)
  425. print(os.linesep.join(stored_files))
  426. @cli.command(name="retrieve")
  427. @click.option('-x', '--context', required=True)
  428. @click.option('-d', '--destination', default=None)
  429. @click.argument('files', nargs=-1)
  430. @click.pass_context
  431. def retrieve_files(ctx, context, destination, files):
  432. """
  433. Retrieve specified files from a <context> bucket
  434. """
  435. ctx.obj['CONTEXT'] = context
  436. s3_bucket, s3 = prep_s3(ctx)
  437. retrieved_files: List[str] = []
  438. if ctx.obj['READ_STDIN']:
  439. files = get_clean_stdin_iterator(click.get_text_stream('stdin'))
  440. for file in files:
  441. file_identity = get_file_identity(ctx.obj, file)
  442. file_destination = file
  443. if destination is not None:
  444. file_destination = os.path.join(destination, file_identity)
  445. try:
  446. s3.fget_object(
  447. s3_bucket,
  448. file_identity,
  449. file_destination
  450. )
  451. retrieved_files.append(file_destination)
  452. except S3Error as e:
  453. if e.code == "NoSuchKey":
  454. LOG.error(f'ERROR: {file_identity} {file_destination} {e}', file=sys.stderr)
  455. else:
  456. LOG.error(e)
  457. except InvalidResponseError as e:
  458. LOG.error(f'ERROR: {file_destination} {e}', file=sys.stderr)
  459. print(os.linesep.join(retrieved_files))
  460. @cli.command(name="clean")
  461. @click.option('-x', '--context', required=True)
  462. @click.option('-d', '--context-data', default=None)
  463. @click.option('-n', '--dry-run/--no-dry-run', default=False)
  464. @click.argument('files', nargs=-1)
  465. @click.pass_context
  466. def clean_files(ctx, context, context_data, dry_run, files):
  467. """
  468. Remove non matching specified files in a <context> bucket for retrieval.
  469. """
  470. ctx.obj['CONTEXT'] = context
  471. s3_bucket, s3 = prep_s3(ctx)
  472. s3_data_bucket = context_data
  473. found_files: List[str] = []
  474. found_data_files: List[str] = []
  475. removed_files: List[str] = []
  476. if ctx.obj['READ_STDIN']:
  477. files = get_clean_stdin_iterator(click.get_text_stream('stdin'))
  478. # Go through and find all matching files
  479. for file in files:
  480. file_identity = f'{get_file_identity(ctx.obj, file)}.json'
  481. try:
  482. if s3_data_bucket is not None:
  483. file_object = s3.get_object(s3_bucket, file_identity)
  484. stored_data = json.load(file_object)
  485. stored_data_file_identity = stored_data['storedAssetIdentity']
  486. found_files.append(file_identity)
  487. found_data_files.append(stored_data_file_identity)
  488. else:
  489. file_object = s3.get_object(s3_bucket, file_identity)
  490. found_files.append(file_identity)
  491. except InvalidResponseError as e:
  492. LOG.error(f'ERROR: InvalidResponseError {file_identity} {e}', file=sys.stderr)
  493. except S3Error as e:
  494. if e.code == "NoSuchKey":
  495. LOG.error(f'ERROR: NoSuchKey {file_identity}', file=sys.stderr)
  496. else:
  497. LOG.error(e)
  498. # print(os.linesep.join(found_objects))
  499. # print(os.linesep.join(found_objects))
  500. found_files = set(found_files)
  501. found_data_files = set(found_data_files)
  502. # Find all objects in s3 bucket
  503. found_objects: List[str] = []
  504. for obj in s3.list_objects_v2(s3_bucket, recursive=False):
  505. if obj.is_dir:
  506. found_objects.extend(list_s3_dir(s3, s3_bucket, obj.object_name))
  507. else:
  508. found_objects.append(obj.object_name)
  509. # print(os.linesep.join(found_objects))
  510. found_data_objects: List[str] = []
  511. for obj in s3.list_objects_v2(s3_data_bucket, recursive=False):
  512. if obj.is_dir:
  513. found_data_objects.extend(list_s3_dir(
  514. s3, s3_data_bucket, obj.object_name))
  515. else:
  516. found_data_objects.append(obj.object_name)
  517. # print(os.linesep.join(found_data_objects))
  518. for file_identity in found_objects:
  519. if file_identity not in found_files:
  520. if dry_run:
  521. removed_files.append(f'{s3_bucket}:{file_identity}')
  522. else:
  523. try:
  524. s3.remove_object(s3_bucket, file_identity)
  525. removed_files.append(f'{s3_bucket}:{file_identity}')
  526. except InvalidResponseError as e:
  527. LOG.error(
  528. f'ERROR: {s3_bucket}:{file_identity} {e}', file=sys.stderr)
  529. for file_data_identity in found_data_objects:
  530. if file_data_identity not in found_data_files:
  531. if dry_run:
  532. removed_files.append(f'{s3_data_bucket}:{file_data_identity}')
  533. else:
  534. try:
  535. s3.remove_object(s3_data_bucket, file_data_identity)
  536. removed_files.append(
  537. f'{s3_data_bucket}:{file_data_identity}')
  538. except InvalidResponseError as e:
  539. LOG.error(
  540. f'ERROR: {s3_data_bucket}:{file_data_identity} {e}', file=sys.stderr)
  541. print(os.linesep.join(removed_files))
  542. ######################################
  543. # Asset Compression Focused Commands #
  544. ######################################
  545. @cli.command(name="compress")
  546. @click.option('-p', '--profile', default='default')
  547. @click.option('-c', '--content', default='all')
  548. @click.option('-d', '--destination', default=None)
  549. @click.option('--print-input-and-identity/--no-print-input-and-identity', default=False)
  550. @click.argument('files', nargs=-1)
  551. @click.pass_context
  552. def compress_assets(ctx, profile, content, destination, print_input_and_identity, files):
  553. """
  554. Compress the request files and store them in a storage bucket.
  555. """
  556. profiles = ctx.obj['CONFIG']['profiles']
  557. if profile not in profiles:
  558. raise ValueError(f'Unrecognized profile: {profile}')
  559. default_profile: Dict[str, any] = profiles['default']
  560. profile: Dict[str, any] = profiles[profile]
  561. if content != 'all':
  562. if content not in profile and content not in default_profile:
  563. raise ValueError(f'Unrecognized content: {content}')
  564. content_configurations = []
  565. if content == 'all':
  566. content_names: set = set()
  567. for content_name in profile.keys():
  568. content_names.add(content_name)
  569. content_configurations.append(profile[content_name])
  570. for content_name in default_profile.keys():
  571. if content_name not in content_names:
  572. content_names.add(content_name)
  573. content_configurations.append(default_profile[content_name])
  574. else:
  575. if content in profile:
  576. content_configurations.append(profile[content])
  577. else:
  578. content_configurations.append(default_profile[content])
  579. if ctx.obj['READ_STDIN']:
  580. files = get_clean_stdin_iterator(click.get_text_stream('stdin'))
  581. if destination is None:
  582. destination = tempfile.mkdtemp()
  583. task_output = []
  584. tasks = []
  585. follow_up_tasks = []
  586. def store_filename(storage_list: List[str], filename: str):
  587. """
  588. A simple lambda wrapper to asynchronously add processed files to the list
  589. :param storage_list:
  590. :param filename:
  591. :return:
  592. """
  593. return lambda: storage_list.append(filename)
  594. def queue_follow_up_task_if_keep_smaller_input(follow_up_tasks, input_file: str, output_file: str, keep_smaller_input: bool = True):
  595. """
  596. A lambda wrapper that handles keeping the smallest of the two files.
  597. """
  598. if keep_smaller_input:
  599. command = f"cp {input_file} {output_file}"
  600. def task():
  601. input_size = os.path.getsize(input_file)
  602. output_size = os.path.getsize(output_file)
  603. if output_size > input_size:
  604. follow_up_tasks.append(
  605. run_command_shell(
  606. command,
  607. stdout=asyncio.subprocess.DEVNULL,
  608. stderr=asyncio.subprocess.DEVNULL,
  609. on_success=[store_filename(
  610. task_output,
  611. f'Preserved smaller "{input_file}" {output_size} > {input_size}'
  612. )]
  613. )
  614. )
  615. return task
  616. return lambda: True
  617. for input_file in files:
  618. for content_configuration in content_configurations:
  619. if any([input_file.endswith(extension) for extension in content_configuration['extensions']]):
  620. file = input_file
  621. file_extension = pathlib.Path(input_file).suffix
  622. if 'REMOVE_PREFIX' in ctx.obj and ctx.obj['REMOVE_PREFIX'] is not None:
  623. file = strip_prefix(ctx.obj['REMOVE_PREFIX'], input_file)
  624. if 'preserveInputExtension' in content_configuration \
  625. and content_configuration['preserveInputExtension']:
  626. output_file = os.path.join(destination, file)
  627. else:
  628. output_file_without_ext = os.path.splitext(
  629. os.path.join(destination, file))[0]
  630. output_file = f'{output_file_without_ext}.{content_configuration["outputExtension"]}'
  631. output_file_identity = get_file_identity(
  632. {'REMOVE_PREFIX': destination}, output_file)
  633. output_file_dir = os.path.dirname(output_file)
  634. os.makedirs(output_file_dir, exist_ok=True)
  635. if 'preserveSmallerInput' in content_configuration:
  636. preserve_smaller_input = bool(
  637. content_configuration['preserveSmallerInput'])
  638. else:
  639. preserve_smaller_input = True
  640. if 'forcePreserveSmallerInput' in content_configuration:
  641. force_preserve_smaller_input = bool(
  642. content_configuration['forcePreserveSmallerInput'])
  643. else:
  644. force_preserve_smaller_input = False
  645. # Only preserve the input if requested AND the extensions of the input and the output match
  646. preserve_smaller_input = preserve_smaller_input and (
  647. force_preserve_smaller_input or file_extension == content_configuration["outputExtension"])
  648. command: str = content_configuration['command'] \
  649. .replace('{input_file}', f'\'{input_file}\'') \
  650. .replace('{output_file}', f'\'{output_file}\'')
  651. tasks.append(
  652. run_command_shell(
  653. command,
  654. stdout=asyncio.subprocess.DEVNULL,
  655. stderr=asyncio.subprocess.DEVNULL,
  656. on_success=[store_filename(
  657. task_output,
  658. f'{input_file}\t{output_file_identity}' if print_input_and_identity else output_file
  659. ), queue_follow_up_task_if_keep_smaller_input(
  660. follow_up_tasks,
  661. input_file,
  662. output_file,
  663. preserve_smaller_input
  664. )]
  665. )
  666. )
  667. results = run_asyncio_commands(
  668. tasks, max_concurrent_tasks=ctx.obj['CONFIG']['concurrency']
  669. )
  670. follow_up_results = run_asyncio_commands(
  671. follow_up_tasks, max_concurrent_tasks=ctx.obj['CONFIG']['concurrency']
  672. )
  673. print(os.linesep.join(task_output))
  674. if __name__ == '__main__':
  675. cli(obj={})