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([el[0].replace('"', '""'), el[1].replace('"', '""')]) 

42 self.data_to_store = list() 

43 

44 def get_value(self, id_string): 

45 ''' 

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

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

48 ''' 

49 if id_string in self.data: 

50 return set(self.data[id_string]) 

51 

52 def add_value(self, id_string, value): 

53 ''' 

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

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

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

57 ''' 

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

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

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

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

62 

63 def __load_csv(self) -> None: 

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

65 for cur_file in cur_files: 

66 if cur_file.endswith('.csv'): 

67 file_path = cur_dir + sep + cur_file 

68 with open(file_path, 'r', encoding='utf-8') as f: 

69 reader = DictReader(f) 

70 for row in reader: 

71 self.data.setdefault(row['id'], set()) 

72 self.data[row['id']].add(row['value'])