from xml.etree import ElementTree import logging import re import sys import gzip import pathlib from progress_bar import progress from word import Word def is_root_id(id_): return len(id_.split('.')) == 3 def load_files(args, database): filenames = args.input skip_id_check = args.skip_id_check do_msd_translate = not args.no_msd_translate database.init("CREATE TABLE Files ( filename varchar(2048) )") for idx, fname in enumerate(filenames): print("FILE ", fname, "{}/{}".format(idx, len(filenames))) extension = pathlib.Path(fname).suffix # check if file with the same name already loaded... loaded = database.execute("SELECT * FROM Files WHERE filename=?", (fname, )).fetchone() if loaded is not None: print("ALREADY LOADED") continue if extension == ".xml": et = load_xml(fname) yield file_sentence_generator(et, skip_id_check, do_msd_translate, args.pc_tag) elif extension == ".gz": yield load_csv(fname, True) else: yield load_csv(fname, False) # else: # raise NotImplementedError("Unknown file extension: {}".format(extension)) database.execute("INSERT INTO Files (filename) VALUES (?)", (fname, )) database.commit() def lines_gz(filename): with gzip.open(filename, 'r') as fp: for line in progress(fp, 'load-gz'): yield line.decode('utf8') def lines_csv(filename): with open(filename, 'r') as fp: for line in progress(fp, 'load-csv'): yield line def load_csv(filename, compressed): result = [] bad_sentence = False words = {} links = [] def sentence_end(bad_sentence): if bad_sentence: return for lfrom, ldest, ana in links: if lfrom not in words or ldest not in words: logging.warning("Bad link in sentence: " + line_split[0]) continue words[lfrom].add_link(ana, words[ldest]) result.extend(words.values()) line_gen = lines_gz if compressed else lines_csv for line in line_gen(filename): line_str = line.strip() line_fixed = line_str.replace(',', '\t').replace('\t\t\t', '\t,\t') line_split = line_fixed.split("\t") if line_split[1] == "1" and len(words) > 0: sentence_end(bad_sentence) bad_sentence = False links = [] words = {} try: sid, wid, text, msd, lemma, link_src, link_type = line_split except ValueError: bad_sentence = True full_id = "{}.{}".format(sid, wid) words[wid] = Word(lemma, msd, full_id, text, True) if link_src != '0': links.append((link_src, wid, link_type)) sentence_end(bad_sentence) return result def load_xml(filename): with open(filename, 'r') as fp: content = fp.read() xmlstring = re.sub(' xmlns="[^"]+"', '', content, count=1) xmlstring = xmlstring.replace(' xml:', ' ') return ElementTree.XML(xmlstring) def file_sentence_generator(et, skip_id_check, do_msd_translate, pc_tag): words = {} sentences = list(et.iter('s')) for sentence in progress(sentences, "load-text"): for w in sentence.iter("w"): words[w.get('id')] = Word.from_xml(w, do_msd_translate) for pc in sentence.iter(pc_tag): words[pc.get('id')] = Word.pc_word(pc, do_msd_translate) for l in sentence.iter("link"): if 'dep' in l.keys(): ana = l.get('afun') lfrom = l.get('from') dest = l.get('dep') else: ana = l.get('ana') if ana[:4] != 'syn:': # dont bother... continue ana = ana[4:] lfrom, dest = l.get('target').replace('#', '').split() if lfrom in words: if not skip_id_check and is_root_id(lfrom): logging.error("NOO: {}".format(lfrom)) sys.exit(1) if dest in words: next_word = words[dest] words[lfrom].add_link(ana, next_word) else: logging.error("Unknown id: {}".format(dest)) sys.exit(1) else: # strange errors, just skip... pass return list(words.values())