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_WORKER_PATTERN: typing.Final = py_re.compile( 

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

11 flags=settings.TYPICAL_RE_FLAGS, 

12) 

13_BROKER_PATTERNS: typing.Final = types.MappingProxyType( 

14 { 

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

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

17 "postgresql": py_re.compile( 

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

19 flags=settings.TYPICAL_RE_FLAGS, 

20 ), 

21 }, 

22) 

23_QUEUE_IMPORT_PATTERNS: typing.Final = types.MappingProxyType( 

24 { 

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

26 for one_queue in TaskQueueEnum 

27 }, 

28) 

29_EMPTY_FEATURES: typing.Final = TaskQueueFeatures( 

30 queues_used=frozenset(), 

31 has_workers=False, 

32 brokers_detected=frozenset(), 

33) 

34 

35 

36def _find_used_queues(raw_source: str, /) -> frozenset[str]: 

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

38 return frozenset( 

39 one_queue.value 

40 for one_queue, one_pattern in _QUEUE_IMPORT_PATTERNS.items() 

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

42 ) 

43 

44 

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

46 queues_found: typing.Final = _find_used_queues(raw_source) 

47 if not queues_found: 

48 return _EMPTY_FEATURES 

49 return TaskQueueFeatures( 

50 queues_used=queues_found, 

51 has_workers=bool(_WORKER_PATTERN.search(raw_source)), 

52 brokers_detected=frozenset( 

53 one_broker_name 

54 for one_broker_name, broker_pattern in _BROKER_PATTERNS.items() 

55 if broker_pattern.search(raw_source) 

56 ), 

57 )