Coverage for src / anpr2mqtt / tools.py: 82%

1import json 

2from pathlib import Path 

3from typing import TYPE_CHECKING, Any, Literal 

4 

5import structlog 

6from PIL import Image 

7from pydantic import BaseModel, Field 

8from pydantic_settings import ( 

9 BaseSettings, 

10 CliApp, 

11 CliPositionalArg, 

12 CliSubCommand, 

13 SettingsConfigDict, 

14) 

15 

16from anpr2mqtt.api_client import DVLAClient 

17from anpr2mqtt.event_handler import examine_file, scan_ocr_fields 

18from anpr2mqtt.settings import DVLASettings, EventSettings, OCRFieldSettings, OCRSettings 

19 

20if TYPE_CHECKING: 

21 from anpr2mqtt.const import ImageInfo 

22 

23log = structlog.get_logger() 

24 

25 

26class OCRTool(BaseModel): 

27 event: EventSettings | None = EventSettings() 

28 ocr: OCRFieldSettings | None = None 

29 log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO" 

30 image_file: CliPositionalArg[str] 

31 

32 def cli_cmd(self) -> None: 

33 structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(self.log_level)) 

34 

35 image_path: Path = self.event.watch_path if self.event is not None else Path() 

36 image: Image.Image | None = Image.open(image_path / self.image_file) 

37 if self.ocr is not None: 

38 ocr_settings: OCRSettings = OCRSettings(fields={"hik_direction": self.ocr}) 

39 else: 

40 ocr_settings = OCRSettings() 

41 log.debug("ocr_files: ocr_settings->%s", ocr_settings) 

42 if image and self.event: 

43 print(scan_ocr_fields(image, self.event, ocr_settings)) # noqa: T201 

44 else: 

45 print("Image can't be loaded") # noqa: T201 

46 

47 

48class ListTool(BaseModel): 

49 event: EventSettings = EventSettings() 

50 log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO" 

51 

52 def cli_cmd(self) -> None: 

53 structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(self.log_level)) 

54 log.info(f"list_dir: {self.event.event} from {self.event.watch_path.resolve()}") 

55 for p in self.event.watch_path.iterdir(): # ty:ignore[invalid-argument-type] 

56 results: ImageInfo | None = examine_file(p, self.event.image_name_re) 

57 if results is not None: 

58 print(f"{results.target}: timestamp={results.timestamp},ext={results.ext}") # noqa: T201 

59 

60 

61class DVLATool(BaseModel): 

62 dvla: DVLASettings = DVLASettings() 

63 test: bool = Field(default=False, description="Use DVLA UAT environment") 

64 log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO" 

65 registration: CliPositionalArg[str] 

66 

67 def cli_cmd(self) -> None: 

68 structlog.configure(wrapper_class=structlog.make_filtering_bound_logger(self.log_level)) 

69 if not self.dvla.api_key: 

70 print("Error: DVLA API key required (--dvla.api_key or DVLA__API_KEY env var)") # noqa: T201 

71 return 

72 print( # noqa: T201 

73 "Caching for {} at {}".format(self.dvla.cache_ttl, "in memory" if not self.dvla.cache_dir else self.dvla.cache_dir) 

74 ) 

75 client = DVLAClient( 

76 api_key=self.dvla.api_key, 

77 cache_type=self.dvla.cache_type, 

78 cache_dir=self.dvla.cache_dir, 

79 cache_ttl=self.dvla.cache_ttl, 

80 test=self.test, 

81 ) 

82 result: dict[str, Any] = client.lookup(self.registration.upper()) 

83 print(json.dumps(result, indent=2)) # noqa: T201 

84 

85 

86class Tools(BaseSettings, cli_parse_args=True, cli_exit_on_error=True): 

87 model_config = SettingsConfigDict( 

88 env_nested_delimiter="__", 

89 env_ignore_empty=True, 

90 cli_avoid_json=True, 

91 ) 

92 ocr_file: CliSubCommand[OCRTool] 

93 list_dir: CliSubCommand[ListTool] 

94 dvla_lookup: CliSubCommand[DVLATool] 

95 

96 def cli_cmd(self) -> None: 

97 CliApp.run_subcommand(self) 

98 

99 

100def tools() -> None: 

101 CliApp.run(model_cls=Tools)