debeir.core.parser

  1import abc
  2import csv
  3import dataclasses
  4import json
  5from collections import defaultdict
  6from dataclasses import dataclass
  7from typing import Dict, List
  8from xml.etree import ElementTree as ET
  9
 10import dill
 11import pandas as pd
 12
 13
 14# TODO: Parse fields can come from a config or ID_fields
 15# TODO: move _get_topics to private cls method with arguments, and expose get_topics as an instance method.
 16
 17
 18@dataclass(init=True)
 19class Parser:
 20    """
 21    Parser interface
 22    """
 23
 24    id_field: object
 25    parse_fields: List[str]
 26
 27    @classmethod
 28    def normalize(cls, input_dict) -> Dict:
 29        """
 30        Flatten the dictionary, i.e. from Dict[int, Dict] -> Dict[str, str_or_int]
 31
 32        :param input_dict:
 33        :return:
 34        """
 35        return pd.io.json.json_normalize(input_dict,
 36                                         sep=".").to_dict(orient='records')[0]
 37
 38    def get_topics(self, path, *args, **kwargs):
 39        """
 40        Instance method for getting topics, forwards instance self parameters to the _get_topics class method.
 41        """
 42
 43        self_kwargs = vars(self)
 44        kwargs.update(self_kwargs)
 45
 46        return self._get_topics(path, *args, **kwargs)
 47
 48    @classmethod
 49    @abc.abstractmethod
 50    def _get_topics(cls, path, *args, **kwargs) -> Dict[int, Dict[str, str]]:
 51        raise NotImplementedError
 52
 53
 54@dataclasses.dataclass(init=True)
 55class PickleParser(Parser):
 56    """
 57    Load topics from a pickle file
 58    """
 59
 60    @classmethod
 61    def _get_topics(cls, path, *args, **kwargs) -> Dict[int, Dict[str, str]]:
 62        return dill.load(path)
 63
 64
 65@dataclasses.dataclass(init=True)
 66class XMLParser(Parser):
 67    """
 68    Load topics from an XML file
 69    """
 70    topic_field_name: str
 71    id_field: str
 72    parse_fields: List[str]
 73
 74    @classmethod
 75    def _recurse_to_child_node(cls, node: ET.Element, track: List):
 76        """
 77        Helper method to get all children nodes for text extraction in an xml.
 78
 79        :param node: Current node
 80        :param track: List to track nodes
 81        :return:
 82        """
 83        if len(node.getchildren()) > 0:
 84            for child in node.getchildren():
 85                track.append(cls._recurse_to_child_node(child, track))
 86
 87        return node
 88
 89    @classmethod
 90    def unwrap(cls, doc_dict, key):
 91        """
 92        Converts defaultdict to dict and list of size 1 to just the element
 93
 94        :param doc_dict:
 95        :param key:
 96        """
 97        if isinstance(doc_dict[key], defaultdict):
 98            doc_dict[key] = dict(doc_dict[key])
 99
100            for e_key in doc_dict[key]:
101                cls.unwrap(doc_dict[key], e_key)
102
103        if isinstance(doc_dict[key], list):
104            if len(doc_dict[key]) == 1:
105                doc_dict[key] = doc_dict[key][0]
106
107    def _get_topics(self, path, *args, **kwargs) -> Dict[int, Dict[str, str]]:
108        all_topics = ET.parse(path).getroot()
109        qtopics = {}
110
111        for topic in all_topics.findall(self.topic_field_name):
112            _id = topic.attrib[self.id_field]
113            if _id.isnumeric():
114                _id = int(_id)
115
116            if self.parse_fields:
117                temp = {}
118                for field in self.parse_fields:
119                    try:
120                        temp[field] = topic.find(field).text.strip()
121                    except:
122                        continue
123
124                qtopics[_id] = temp
125            else:
126                #  The topic contains the text
127                qtopics[_id] = {"query": topic.text.strip()}
128
129        return qtopics
130
131
132@dataclasses.dataclass
133class CSVParser(Parser):
134    """
135    Loads topics from a CSV file
136    """
137    id_field = "id"
138    parse_fields = ["Text"]
139
140    def __init__(self, id_field=None, parse_fields=None):
141        if parse_fields is None:
142            parse_fields = ["id", "text"]
143
144        if id_field is None:
145            id_field = "id"
146
147        super().__init__(id_field, parse_fields)
148
149    @classmethod
150    def _get_topics(cls, csvfile, dialect="excel",
151                    id_field: str = None,
152                    parse_fields: List[str] = None,
153                    *args, **kwargs) -> Dict[int, Dict[str, str]]:
154        topics = {}
155
156        if isinstance(csvfile, str):
157            csvfile = open(csvfile, 'rt')
158
159        if id_field is None:
160            id_field = cls.id_field
161
162        if parse_fields is None:
163            parse_fields = cls.parse_fields
164
165        reader = csv.DictReader(csvfile, dialect=dialect)
166        for row in reader:
167            temp = {}
168
169            for field in parse_fields:
170                temp[field] = row[field]
171
172            topics[row[id_field]] = temp
173
174        return topics
175
176
177@dataclasses.dataclass(init=True)
178class TSVParser(CSVParser):
179
180    @classmethod
181    def _get_topics(cls, tsvfile, *args, **kwargs) -> Dict[int, Dict[str, str]]:
182        return CSVParser._get_topics(tsvfile, *args, dialect='excel-tab', **kwargs)
183
184
185@dataclasses.dataclass(init=True)
186class JsonLinesParser(Parser):
187    """
188    Loads topics from a jsonl file,
189    a JSON per line
190
191    Provide parse_fields, id_field and whether to ignore full matches on json keys
192    secondary_id appends to the primary id as jsonlines are flattened structure and may contain duplicate ids.
193    """
194    parse_fields: List[str]
195    id_field: str
196    ignore_full_match: bool = True
197    secondary_id: str = None
198
199    @classmethod
200    def _get_topics(cls, jsonlfile, id_field, parse_fields,
201                    ignore_full_match=True, secondary_id=None, *args, **kwargs) -> Dict[str, Dict]:
202        with open(jsonlfile, "r") as jsonl_f:
203            topics = {}
204
205            for jsonl in jsonl_f:
206                json_dict = json.loads(jsonl)
207                _id = json_dict.pop(id_field)
208
209                if secondary_id:
210                    _id = str(_id) + "_" + str(json_dict[secondary_id])
211
212                for key in list(json_dict.keys()):
213                    found = False
214                    for _key in parse_fields:
215                        if ignore_full_match:
216                            if key in _key or key == _key or _key in key:
217                                found = True
218                        else:
219                            if _key == key:
220                                found = True
221                    if not found:
222                        json_dict.pop(key)
223
224                topics[_id] = json_dict
225
226        return topics