Coverage for archdocs/features/task_queues/parser.py: 100%

1import re as py_re 

2import types 

3import typing 

4 

5from archdocs import prefilter, settings 

6from archdocs.features.task_queues.const import TaskQueueEnum, TaskQueueFeatures 

7 

8 

9_TASK_DECORATOR_PATTERNS: typing.Final = py_re.compile( 

10 r"@(?:\w+\.)?(?:task|actor|job)\b", 

11 flags=settings.TYPICAL_RE_FLAGS, 

12) 

13_WORKER_PATTERNS: typing.Final = py_re.compile( 

14 r"\b(?:celery\s+worker|taskiq\s+worker|arq\s+worker|rq\s+worker|dramatiq\s+worker|huey\s+worker)\b", 

15 flags=settings.TYPICAL_RE_FLAGS, 

16) 

17_BROKER_PATTERNS: typing.Final = types.MappingProxyType( 

18 { 

19 "redis": py_re.compile(r"\b(?:redis://|RedisSettings|RedisHuey)\b", flags=settings.TYPICAL_RE_FLAGS), 

20 "rabbitmq": py_re.compile(r"\b(?:amqp://|rabbitmq://|RabbitMQBroker)\b", flags=settings.TYPICAL_RE_FLAGS), 

21 "postgresql": py_re.compile( 

22 r"\b(?:postgres://|postgresql://|PostgreSQLBroker)\b", 

23 flags=settings.TYPICAL_RE_FLAGS, 

24 ), 

25 }, 

26) 

27_QUEUE_IMPORT_PATTERNS: typing.Final = types.MappingProxyType( 

28 { 

29 one_queue: py_re.compile(rf"\b(?:from|import)\s+{one_queue.value}\b", flags=settings.TYPICAL_RE_FLAGS) 

30 for one_queue in TaskQueueEnum 

31 }, 

32) 

33_EMPTY_FEATURES: typing.Final = TaskQueueFeatures( 

34 queues_used=frozenset(), 

35 has_tasks=False, 

36 has_workers=False, 

37 brokers_detected=frozenset(), 

38) 

39 

40 

41def _find_used_queues(raw_source: str, /) -> set[str]: 

42 lowered_source: typing.Final = raw_source.lower() 

43 return { 

44 one_queue.value 

45 for one_queue, one_pattern in _QUEUE_IMPORT_PATTERNS.items() 

46 if prefilter.contains_any_literal(lowered_source, (one_queue.value,)) and one_pattern.search(raw_source) 

47 } 

48 

49 

50def find_task_queue_features(raw_source: str) -> TaskQueueFeatures: 

51 queues_found: typing.Final = _find_used_queues(raw_source) 

52 if not queues_found: 

53 return _EMPTY_FEATURES 

54 return TaskQueueFeatures( 

55 queues_used=frozenset(queues_found), 

56 has_tasks=bool(_TASK_DECORATOR_PATTERNS.search(raw_source)), 

57 has_workers=bool(_WORKER_PATTERNS.search(raw_source)), 

58 brokers_detected=frozenset( 

59 one_broker_name 

60 for one_broker_name, broker_pattern in _BROKER_PATTERNS.items() 

61 if broker_pattern.search(raw_source) 

62 ), 

63 )