1#!/usr/bin/env python3
2
3"""Wrapper around the standard CSV module to:
4
5- Simpler construction arguments
6- Write timestamps in a standard way that attempts to give the best results
7 in Excel.
8 """
9
10
11
12
13import sys
14import csv
15from datetime import datetime
16from enum import Enum
17
18# strftime format for the timestamp column of CSV file
19CSV_TIME_FORMAT = '%Y-%m-%d %H:%M:%S'
20
21CSV_MODE = Enum(str('CSV_MODE'), str('READ WRITE')) # py2
22
23
24def csv_to_datetime(s):
25 """Convert string of format 2015-04-05 12:34:56.123 to datetime.
26 Milliseconds part optional.
27 """
28 if '.' in s:
29 res = datetime.strptime(s, '%Y-%m-%d %H:%M:%S.%f')
30
31 else:
32 res = datetime.strptime(s, '%Y-%m-%d %H:%M:%S')
33
34 return res
35
36
37class CSV:
38 """Thin wrapper around the built in csv object with:
39
40 - a simpler constructor incl. writing to stdout
41 - automatic conversion of datetimes to excel friendly format
42 """
43
44 APPEND = 'a'
45 READ = 'r'
46 WRITE = 'w'
47
48 def __init__(self, path=None, mode=None, headings=None, dictreader=False):
49 """Open a CSV file for reading (`mode`='r') or writing (`mode`='r').
50 `filename` should be a Path object.
51 If filename is empty the CSV will be written to stdout.
52
53 For writing only the headings (first row) can be specified here, otherwise
54 they can be passed in the first writerow() call.
55
56 If mode is 'r' and dict is true, rows are returned as dictionaries.
57 """
58
59 # if created with no path, no mode, assume we write to terminal
60 if path is None and mode is None:
61 mode = CSV.WRITE
62
63 # by default read from the named file
64 elif mode is None:
65 mode = CSV.READ
66
67 self.path = path
68 self.mode = mode
69 self.dictreader = dictreader
70
71 if headings is not None:
72 if mode is None:
73 mode = CSV.WRITE
74
75 elif mode != CSV.WRITE:
76 raise ValueError('CSV headings specified but mode is not w')
77
78 if mode == CSV.READ:
79 self.writer = None
80 self.reader = csv.reader(path.open('r'), dialect='excel')
81 if dictreader:
82 self.reader = csv.DictReader(self.reader)
83
84 elif mode == CSV.WRITE:
85 self.reader = None
86 if path is None:
87 # write to stdout
88 self.writer = None
89
90 else:
91 # without buffering we get weird truncation using fake writes from iasi_csd_ingester
92 # but pathlib refuses to open text files without buffering
93 # self.writer = csv.writer(path.open('w'), dialect='excel')
94 self.writer = csv.writer(path.open('w'), dialect='excel')
95
96 elif mode == CSV.APPEND:
97 self.reader = None
98 if path is None:
99 self.writer = None
100
101 else:
102 filemode = 'a'
103 if sys.version_info.major < 3:
104 filemode += 'b'
105
106 self.writer = csv.writer(path.open(filemode), dialect='excel')
107 # self.writer = csv.writer(path.open('a', buffering=0), dialect='excel')
108
109 else:
110 raise ValueError('Unknown value {m} for mode'.format(m=mode))
111
112 if headings is not None:
113 self.writerow(headings)
114
115 def writerow(self, row):
116 """Write a row to the output file. datetime objects are converted to a general,
117 Excel-friendly format."""
118 out_row = []
119 for item in row:
120 if isinstance(item, datetime):
121 if item.microsecond != 0:
122 out_row.append('{coarse}.{ms:03}'.format(
123 coarse=item.strftime(CSV_TIME_FORMAT),
124 ms=item.microsecond // 1000))
125
126 else:
127 out_row.append('{coarse}'.format(
128 coarse=item.strftime(CSV_TIME_FORMAT)))
129
130 else:
131 out_row.append(item)
132
133 if self.writer is None:
134 print(', '.join(str(i) for i in out_row))
135
136 else:
137 self.writer.writerow(out_row)
138
139 def append(self, row):
140 """Compatibility with prettyprint.Table."""
141
142 self.writerow(row)
143
144 def __iter__(self):
145 return self.reader.__iter__()