Coverage for collectra/commons/base.py: 94.55%

1import enum 

2from abc import ABC, abstractmethod 

3from dataclasses import dataclass, field 

4from pathlib import Path 

5from typing import Generic, TypeVar 

6 

7__all__ = ["T", "BaseEntity", "Node", "NodeStatus", "TaskContext"] 

8 

9T = TypeVar("T") 

10 

11 

12@dataclass 

13class TaskContext: 

14 """Runtime context for task execution. 

15 

16 This provides tasks with runtime information without polluting 

17 their function signatures or requiring global state. 

18 

19 Attributes: 

20 usage_file: Optional path to write LLM usage statistics. 

21 render: Whether to render workflow visualizations. 

22 file_path: Current file being processed. 

23 single_run: Whether to stop after first task execution. 

24 """ 

25 

26 usage_file: Path | None = None 

27 render: bool = False 

28 file_path: Path | None = None 

29 single_run: bool = False 

30 

31 

32class BaseEntity(ABC, Generic[T]): 

33 

34 def __init__(self, name: str) -> None: 

35 self.name = name 

36 

37 def __str__(self) -> str: 

38 return f"{self.name}\n{self.get_class_path()}" 

39 

40 @abstractmethod 

41 def __call__(self) -> T: 

42 pass 

43 

44 def get_name(self) -> str: 

45 """Get the name of the task. 

46 

47 Returns: 

48 str: The name of the task. 

49 """ 

50 return self.name 

51 

52 def serialize(self) -> dict: 

53 serialized = dict(type=self.get_class_path()) 

54 for key, value in self.attributes.items(): 

55 if isinstance(value, Path): 

56 value = value.name 

57 serialized[key] = value 

58 return serialized 

59 

60 @property 

61 def attributes(self) -> dict: 

62 ignore = self.attributes_to_ignore() 

63 return {k: v for k, v in self.__dict__.items() if k not in ignore} 

64 

65 def attributes_to_ignore(self) -> set: 

66 return set() 

67 

68 @classmethod 

69 def get_class_path(cls: type) -> str: 

70 """Return the fully qualified path for a class. 

71 

72 Example: 

73 >>> from collectra.images import Image 

74 >>> get_class_path(Image) 

75 'collectra.images.Image' 

76 """ 

77 return f"{cls.__module__.split('.')[0]}.{cls.__name__}" 

78 

79 

80class NodeStatus(enum.Enum): 

81 NOT_READY = "not_ready" 

82 READY = "ready" 

83 RUNNING = "running" 

84 COMPLETED = "completed" 

85 FAILED = "failed" 

86 

87 def __str__(self) -> str: 

88 return self.value 

89 

90 

91@dataclass 

92class Node: 

93 

94 name: str 

95 status: NodeStatus = field(init=False, default=NodeStatus.NOT_READY) 

96 

97 def __post_init__(self) -> None: 

98 pass 

99 

100 def process(self) -> list: 

101 return []