-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
executable file
·94 lines (80 loc) · 3.42 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
#!/usr/bin/env python
# vim: set fileencoding=utf-8
import argparse
import cPickle
import functools
import logging
import logging.config
import os.path
import sys
import sqlite3
from gzip import GzipFile
from bz2 import BZ2File
from lzma import LZMAFile
from bac2010parser.rezultate.elev import Elev
from bac2010parser import get_data_from_file
from bac2010parser.rezultate import get_elev_from_mainTable
COMPRESSED_FILE_CLASSES = {'.gz': GzipFile,
'.bz2': BZ2File,
'.xz': LZMAFile}
def open_compressed_file(filename):
global COMPRESSED_FILE_CLASSES
if filename == '-':
return sys.stdin
ext = os.path.splitext(filename)[1]
f = COMPRESSED_FILE_CLASSES.get(ext, file)(filename)
dir(f) # workaround for https://bugzilla.redhat.com/show_bug.cgi?id=720111
return f
def get_data(filenames):
for filename in filenames:
logging.info("Extracting from %s" % (filename,))
with open_compressed_file(filename) as f:
for i in get_data_from_file(f, get_elev_from_mainTable):
yield i
def output_python(f, record):
f.write(repr(record))
f.write('\n#######################################################################\n')
def output_pickle(f, record):
cPickle.dump(record, f, -1)
def db_create_table(conn, tbl_name):
f = ('%s varchar(100)' % (i,) for i in Elev._fields)
conn.execute('CREATE TABLE %s (%s)' % (tbl_name, ",".join(f)))
def parse_args():
parser = argparse.ArgumentParser(
description=u'Extrage informații despre elevi din fișiere HTML')
parser.add_argument('filenames', metavar='FILE', type=str, nargs='+',
help=u'O pagină de pe site-ul edu.ro. Folosiți - pentru stdin.')
parser.add_argument('-o', '--output', metavar='OUTPUT',
type=str, default=sys.stdout,
help=u'Fișierul de ieșire. Implicit e stdout.')
parser.add_argument('--format', metavar='FORMAT',
type=str, choices=('python', 'pickle', 'sqlite'),
default='python',
help=u'''Formatul de ieșire. Formate suportate: %(choices)s. Format implicit: %(default)s.''')
parser.add_argument('--dbtable', metavar='DBTABLE',
type=str, default='rezultate',
help=u'Numele tabelului din baza de date. Nume implicit: %(default)s.')
args = parser.parse_args()
return args
def main():
args = parse_args()
if args.format == 'python' or args.format == 'pickle':
if args.format == 'python':
fout = open(args.output, 'wt') if isinstance(args.output,basestring) else args.output
output = functools.partial(output_python, fout)
elif args.format == 'pickle':
fout = open(args.output, 'wb') if isinstance(args.output,basestring) else args.output
output = functools.partial(output_pickle, fout)
with fout:
for i in get_data(args.filenames):
output(i)
elif args.format == 'sqlite':
con = sqlite3.connect(args.output)
db_create_table(con, args.dbtable)
insert_query = 'INSERT INTO %s VALUES(%s)' % \
(args.dbtable, ','.join('?'*len(Elev._fields)))
with con:
con.executemany(insert_query, get_data(args.filenames))
if __name__ == '__main__':
logging.config.fileConfig('logging.ini')
main()