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

1from pathlib import Path 

2from typing import TYPE_CHECKING, Literal 

3 

4import structlog 

5from PIL import Image 

6from pydantic import BaseModel 

7from pydantic_settings import ( 

8 BaseSettings, 

9 CliApp, 

10 CliPositionalArg, 

11 CliSubCommand, 

12 SettingsConfigDict, 

13) 

14 

15from anpr2mqtt.event_handler import examine_file, scan_ocr_fields 

16from anpr2mqtt.settings import EventSettings, OCRFieldSettings, OCRSettings 

17 

18if TYPE_CHECKING: 

19 from anpr2mqtt.const import ImageInfo 

20 

21log = structlog.get_logger() 

22 

23 

24class OCRTool(BaseModel): 

25 event: EventSettings | None = EventSettings() 

26 ocr: OCRFieldSettings | None = None 

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

28 image_file: CliPositionalArg[str] 

29 

30 def cli_cmd(self) -> None: 

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

32 

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

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

35 if self.ocr is not None: 

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

37 else: 

38 ocr_settings = OCRSettings() 

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

40 if image and self.event: 

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

42 else: 

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

44 

45 

46class ListTool(BaseModel): 

47 event: EventSettings = EventSettings() 

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

49 

50 def cli_cmd(self) -> None: 

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

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

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

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

55 if results is not None: 

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

57 

58 

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

60 model_config = SettingsConfigDict( 

61 env_nested_delimiter="__", 

62 env_ignore_empty=True, 

63 cli_avoid_json=True, 

64 ) 

65 ocr_file: CliSubCommand[OCRTool] 

66 list_dir: CliSubCommand[ListTool] 

67 

68 def cli_cmd(self) -> None: 

69 CliApp.run_subcommand(self) 

70 

71 

72def tools() -> None: 

73 CliApp.run(model_cls=Tools)