Coverage for oc_meta / lib / csvmanager.py: 94%

1#!/usr/bin/python 

2 

3# SPDX-FileCopyrightText: 2019 Silvio Peroni <essepuntato@gmail.com> 

4# SPDX-FileCopyrightText: 2022-2026 Arcangelo Massari <arcangelo.massari@unibo.it> 

5# 

6# SPDX-License-Identifier: ISC 

7 

8from csv import DictReader, writer 

9from os import mkdir, sep, walk 

10from os.path import exists, join 

11from typing import Dict 

12 

13 

14class CSVManager(object): 

15 def __init__(self, output_path: str | None = None): 

16 self._output_path = output_path 

17 self.data: Dict[str, set] = {} 

18 self.data_to_store: list[list[str]] = [] 

19 if output_path is not None: 

20 self.__init_output_dir() 

21 self.__load_csv() 

22 

23 @property 

24 def output_path(self) -> str: 

25 if self._output_path is None: 

26 raise ValueError("output_path is not set") 

27 return self._output_path 

28 

29 def __init_output_dir(self) -> None: 

30 if not exists(self.output_path): 

31 mkdir(self.output_path) 

32 

33 def dump_data(self, file_name: str) -> None: 

34 path = join(self.output_path, file_name) 

35 if not exists(path): 

36 with open(path, "w", encoding="utf-8", newline="") as f: 

37 f.write('"id","value"\n') 

38 with open(path, "a", encoding="utf-8", newline="") as f: 

39 csv_writer = writer(f, delimiter=",") 

40 for el in self.data_to_store: 

41 csv_writer.writerow( 

42 [el[0].replace('"', '""'), el[1].replace('"', '""')] 

43 ) 

44 self.data_to_store = list() 

45 

46 def get_value(self, id_string): 

47 """ 

48 It returns the set of values associated to the input 'id_string', 

49 or None if 'id_string' is not included in the CSV. 

50 """ 

51 if id_string in self.data: 

52 return set(self.data[id_string]) 

53 

54 def add_value(self, id_string, value): 

55 """ 

56 It adds the value specified in the set of values associated to 'id_string'. 

57 If the object was created with the option of storing also the data in a CSV 

58 ('store_new' = True, default behaviour), then it also add new data in the CSV. 

59 """ 

60 self.data.setdefault(id_string, set()) 

61 if value not in self.data[id_string]: 

62 self.data[id_string].add(value) 

63 self.data_to_store.append([id_string, value]) 

64 

65 def __load_csv(self) -> None: 

66 for cur_dir, _, cur_files in walk(self.output_path): 

67 for cur_file in cur_files: 

68 if cur_file.endswith(".csv"): 

69 file_path = cur_dir + sep + cur_file 

70 with open(file_path, "r", encoding="utf-8") as f: 

71 reader = DictReader(f) 

72 for row in reader: 

73 self.data.setdefault(row["id"], set()) 

74 self.data[row["id"]].add(row["value"])