secator 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

secator/runners/_helpers.py

@@ -1,32 +1,28 @@
 import os
 
-import kombu
-import kombu.exceptions
-
 from secator.utils import deduplicate
-from secator.rich import console
 
 
-def run_extractors(results, opts, targets=[]):
+def run_extractors(results, opts, inputs=[]):
 	"""Run extractors and merge extracted values with option dict.
 
 	Args:
 		results (list): List of results.
 		opts (dict): Options.
-		targets (list): Original targets.
+		inputs (list): Original inputs.
 
 	Returns:
-		tuple: targets, options.
+		tuple: inputs, options.
 	"""
 	extractors = {k: v for k, v in opts.items() if k.endswith('_')}
 	for key, val in extractors.items():
 		key = key.rstrip('_')
 		values = extract_from_results(results, val)
 		if key == 'targets':
-			targets = deduplicate(values)
+			inputs = deduplicate(values)
 		else:
 			opts[key] = deduplicate(values)
-	return targets, opts
+	return inputs, opts
 
 
 def extract_from_results(results, extractors):
@@ -73,88 +69,6 @@ def process_extractor(results, extractor, ctx={}):
 	return items
 
 
-def get_task_ids(result, ids=[]):
-	"""Get all Celery task ids recursively.
-
-	Args:
-		result (Union[AsyncResult, GroupResult]): Celery result object.
-		ids (list): List of ids.
-	"""
-	from celery.result import AsyncResult, GroupResult
-	if result is None:
-		return
-
-	try:
-		if isinstance(result, GroupResult):
-			get_task_ids(result.parent, ids=ids)
-
-		elif isinstance(result, AsyncResult):
-			if result.id not in ids:
-				ids.append(result.id)
-
-		if hasattr(result, 'children') and result.children:
-			for child in result.children:
-				get_task_ids(child, ids=ids)
-
-		# Browse parent
-		if hasattr(result, 'parent') and result.parent:
-			get_task_ids(result.parent, ids=ids)
-	except kombu.exceptions.DecodeError as e:
-		console.print(f'[bold red]{str(e)}. Aborting get_task_ids.[/]')
-		return
-
-
-def get_task_data(task_id):
-	"""Get task info.
-
-	Args:
-		task_id (str): Celery task id.
-
-	Returns:
-		dict: Task info (id, name, state, results, chunk_info, count, error, ready).
-	"""
-	from celery.result import AsyncResult
-	res = AsyncResult(task_id)
-	if not res:
-		return
-	try:
-		args = res.args
-		info = res.info
-		state = res.state
-	except kombu.exceptions.DecodeError as e:
-		console.print(f'[bold red]{str(e)}. Aborting get_task_data.[/]')
-		return
-	if not (args and len(args) > 1):
-		return
-	task_name = args[1]
-	data = {
-		'id': task_id,
-		'name': task_name,
-		'state': state,
-		'chunk_info': '',
-		'count': 0,
-		'error': None,
-		'ready': False,
-		'descr': '',
-		'progress': 0,
-		'results': []
-	}
-
-	# Set ready flag
-	if state in ['FAILURE', 'SUCCESS', 'REVOKED']:
-		data['ready'] = True
-
-	# Set task data
-	if info and not isinstance(info, list):
-		data.update(info)
-		chunk = data.get('chunk')
-		chunk_count = data.get('chunk_count')
-		if chunk and chunk_count:
-			data['chunk_info'] = f'{chunk}/{chunk_count}'
-		data['descr'] = data.pop('description', '')
-	return data
-
-
 def get_task_folder_id(path):
 	names = []
 	if not os.path.exists(path):

secator/runners/celery.py

@@ -0,0 +1,18 @@
+from secator.celery_utils import CeleryData
+from secator.runners import Runner
+
+
+class Celery(Runner):
+	def yielder(self):
+		if not self.celery_result:
+			result = self.build_celery_workflow()
+		if self.sync:
+			yield from result.apply().get()
+		yield from CeleryData.iter_results(
+			self.celery_result,
+			ids_map=self.celery_ids_map,
+			print_remote_info=False
+		)
+
+	def error_handler(self, e):
+		self.stop_celery_tasks()