forked from CityofSantaMonica/mds-provider
-
Notifications
You must be signed in to change notification settings - Fork 0
/
files.py
661 lines (521 loc) · 24 KB
/
files.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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
"""
Work with MDS Provider data in JSON files.
"""
import datetime
import hashlib
import json
import pathlib
import urllib
import requests
import pandas as pd
from .encoding import JsonEncoder, TimestampDecoder, TimestampEncoder
from .providers import Provider
from .schemas import SCHEMA_TYPES, STATUS_CHANGES, TRIPS, EVENTS, VEHICLES, Schema
from .versions import UnexpectedVersionError, Version
class BaseFile():
"""
Base class for working with Provider files.
"""
def __init__(self, *sources, **kwargs):
"""
Parameters:
sources: str, Path, list, optional
Zero or more paths to track.
"""
self._sources = []
for source in sources:
if isinstance(source, list):
self._sources.extend([self._parse(s) for s in source])
else:
self._sources.append(self._parse(source))
self._sources = list(filter(None, self._sources))
@property
def file_sources(self):
"""
True if this instance references one or more valid file sources.
"""
return all([self._isfile(s) or self._isurl(s) for s in self._sources])
@classmethod
def _isdir(cls, source):
"""
Return True if source is a valid directory that exists.
"""
path = pathlib.Path(source.path)
return not cls._isfile(source) and path.is_dir() and path.exists()
@classmethod
def _isfile(cls, source):
"""
Return True if path is a valid file that exists.
"""
path = pathlib.Path(source.path)
return not cls._isurl(source) and path.is_file() and path.exists()
@classmethod
def _isurl(cls, source):
"""
Return True if source is a valid URL.
"""
return source.scheme in ("http", "https") and source.netloc
@classmethod
def _parse(cls, source):
"""
Parse a data file source argument into an urllib.parse.ParseResult instance.
"""
return urllib.parse.urlparse(str(source)) if source else None
class ConfigFile(BaseFile):
"""
Work with Provider configuration data in JSON files.
"""
def __init__(self, path=None, provider=None, **kwargs):
"""
Parameters:
path: str, Path, optional
A path to a configuration file.
provider: str, UUID, Provider, optional
An identifier (name, id) for a provider; or a Provider instance. Used to key
configuration data in a dict.
"""
super().__init__(path, **kwargs)
self._config_path = None
# did we get a single file path or a provider?
if len(self._sources) == 1 and self._isfile(self._sources[0]):
self._config_path = pathlib.Path(self._sources[0].path)
# read from the config file
if self._config_path:
config = json.load(self._config_path.open())
search = []
# case-insensitive search in config
if isinstance(provider, Provider):
search.extend([
str(provider.provider_id),
provider.provider_name, provider.provider_name.capitalize(),
provider.provider_name.lower(), provider.provider_name.upper()
])
elif provider:
search.extend([
provider, provider.lower(), provider.capitalize(), provider.upper()
])
for s in set(search):
if s in config:
config = config.pop(s)
break
for k,v in config.items():
setattr(self, k, v)
# set default attributes
else:
defaults = [
("auth_type", "Bearer"),
("headers", {}),
("version", Version.mds_lower()),
("mds_api_suffix", None)
]
for _field, _default in defaults:
setattr(self, _field, _default)
# finally, set from keyword args
for k,v in kwargs.items():
setattr(self, k, v)
def __repr__(self):
return f"<mds.files.ConfigFile ('{self._config_path}')>"
def dump(self, path=None, provider=None, **kwargs):
"""
Convert this instance back into a configuration dict.
Parameters:
path: str, Path, optional
The path to write the configuration data.
provider: str, UUID, Provider, optional
An identifier (name, id) for a provider; or a Provider instance. Used to key
configuration data in a dict.
Additional keyword arguments are passed-through to json.dump().
Return:
dict
With no path information, return a dict of configuration.
ConfigFile
With path information, dump configuration to file path and return this instance.
"""
dump = vars(self)
if provider:
if isinstance(provider, Provider):
provider = provider.provider_name
dump = dict([(provider, dump)])
if path:
json.dump(dump, pathlib.Path(path).open("w"), cls=JsonEncoder, **kwargs)
return self
return dump
class DataFile(BaseFile):
"""
Work with Provider payload data in JSON files.
"""
def __init__(self, record_type=None, *sources, **kwargs):
"""
Parameters:
record_type: str, optional
The type of MDS Provider record to use by default.
sources: str, Path, list, optional
One or more paths to (directories containing) MDS payload (JSON) files to read by default.
Directories are expanded such that all corresponding files within are read.
file_name: str, callable(record_type=str, payloads=list, extension=str, [payload=dict]): str, optional
A str name for the file; or a function receiving record_type, list of payloads,
file extension, and optionally a single payload being written, and returns the str
name for the file.
ls: callable(sources=list): list, optional
A function that receives a list of urllib.parse.ParseResult, and returns the
complete list of file Path objects and URL str to be read.
"""
super().__init__(*sources, **kwargs)
self.record_type = None
if record_type:
if record_type in SCHEMA_TYPES:
self.record_type = record_type
else:
self._sources.append(self._parse(record_type))
file_name = kwargs.get("file_name", self._filename)
if isinstance(file_name, str):
self.file_name = lambda **kwargs: file_name
else:
self.file_name = file_name
self.ls = kwargs.get("ls", self._ls)
def __repr__(self):
return "".join((
"<mds.files.DataFile (",
", ".join([f"'{s}'" for s in [self.record_type]]),
")>"
))
def _default_dir(self):
"""
Get a default Path object for dumping data files.
"""
dirs = [s.path for s in self._sources if self._isdir(s)]
return pathlib.Path(dirs[0]) if len(dirs) == 1 else pathlib.Path(".")
def _record_type_or_raise(self, record_type):
"""
Get a valid record_type or raise an exception.
"""
record_type = record_type or self.record_type
if record_type in SCHEMA_TYPES:
return record_type
raise ValueError(f"A valid record type must be specified. Got {record_type}")
def dump_payloads(self, record_type=None, *payloads, **kwargs):
"""
Write MDS Provider payloads to JSON files.
Parameters:
record_type: str, optional
The type of MDS Provider record.
payloads: dict, iterable
One or more MDS Provider payload dicts to write.
output_dir: str, Path, optional
The directory to write the files.
If this instance was initialized with a single directory source, use that by default.
Otherwise, use the current directory by default.
file_name: str, callable(record_type=str, payloads=list, extension=str, [payload=dict]): str, optional
A str name for the file; or a function receiving record_type, list of payloads,
file extension, and optionally a single payload being written, and returns the str
name for the file.
single_file: bool, optional
True (default) to write the payloads to a single file using the appropriate data structure.
False to write each payload as a dict to its own file.
Additional keyword arguments are passed through to json.dump().
Return:
Path
With single_file=True, the Path object pointing to the file that was written.
With single_file=False, the Path object pointing to the directory where files were written.
None if no files were written.
"""
sources = []
# marker indicates if the original incoming source was just a single dict
dict_source = False
# not a true record_type, but a data source
if record_type and record_type not in SCHEMA_TYPES:
if isinstance(record_type, dict):
sources.append(record_type)
dict_source = True
elif isinstance(record_type, list):
sources.extend(record_type)
elif isinstance(record_type, tuple):
sources.extend(list(record_type))
record_type = None
record_type = record_type or self.record_type
# convert payloads to a flat list of dicts
if isinstance(payloads, tuple) and len(payloads) == 1:
payloads = payloads[0]
if isinstance(payloads, dict):
payloads = [payloads]
dict_source = True
if not isinstance(payloads, list):
payloads = list(payloads)
sources.extend(payloads)
# filter payloads with non-matching record_type
if record_type in SCHEMA_TYPES:
data_key = Schema(record_type).data_key
sources = [p for p in sources if data_key in p["data"]]
if len(sources) == 0:
return None
output_dir = pathlib.Path(kwargs.pop("output_dir", self._default_dir()))
single_file = kwargs.pop("single_file", True)
file_name = kwargs.pop("file_name", self.file_name)
if isinstance(file_name, str):
orig_file_name = file_name
file_name = lambda **kwargs: orig_file_name
output_dir.mkdir(parents=True, exist_ok=True)
if single_file:
version = sources[0]["version"]
encoder = JsonEncoder(date_format="unix", version=version, **kwargs)
# generate a file name for the list of payloads
fname = file_name(record_type=record_type, payloads=sources, extension=".json")
path = pathlib.Path(output_dir, fname)
# dump the single payload or a list of payloads
if dict_source and len(sources) == 1:
path.write_text(encoder.encode(sources[0]))
else:
path.write_text(encoder.encode(sources))
return path
# multi-file
for payload in sources:
version = payload["version"]
encoder = JsonEncoder(date_format="unix", version=version, **kwargs)
# generate a file name for this payload
fname = file_name(record_type=record_type, payloads=sources, extension=".json", payload=payload)
path = pathlib.Path(output_dir, fname)
if sources.index(payload) > 0 and path.exists():
# increment the file number
n = str(sources.index(payload))
# pad with number of zeros based on how many items in the list
nz = len(str(len(sources)))
path = pathlib.Path(str(path).replace(".json", f"_{n.zfill(nz)}.json"))
# dump the payload dict
path.write_text(encoder.encode(payload))
return output_dir
def load_dataframe(self, record_type=None, *sources, **kwargs):
"""
Reads the contents of MDS payload files into tuples of (Version, DataFrame).
Parameters:
record_type: str, optional
The type of MDS Provider record.
sources: str, list, optional
One or more paths to (directories containing) MDS payload (JSON) files.
Directories are expanded such that all corresponding files within are read.
flatten: bool, optional
True (default) to flatten the final result from all sources into a single tuple.
False to keep each result separate.
headers: dict, optional
A dict of headers to send with requests made to URL paths.
Could also be a dict mapping an URL path to headers for that path.
ls: callable(sources=list): list, optional
A function that receives a list of urllib.parse.ParseResult, and returns the
complete list of file Path objects and URL str to be read.
Raise:
UnexpectedVersionError
When flatten=True and a version mismatch is found amongst the data.
ValueError
When neither record_type or instance.record_type is specified.
Return:
tuple (Version, DataFrame)
With flatten=True, a (Version, DataFrame) tuple.
list
With flatten=False, a list of (Version, DataFrame) tuples with length equal to the
total number of payloads across all sources.
"""
record_type = self._record_type_or_raise(record_type)
flatten = kwargs.pop("flatten", True)
# obtain unmodified records
kwargs["flatten"] = False
records = self.load_records(record_type, *sources, **kwargs)
if len(records) == 0:
return records
version = Version(records[0][0])
if flatten:
if not all([Version(v) == version for v,_ in records]):
unexpected = [Version(v) for v,_ in records if Version(v) != version][0]
raise UnexpectedVersionError(unexpected, version)
# combine each record list
records = [item for _,data in records for item in data]
return version, pd.DataFrame.from_records(records)
else:
# list of version, DataFrame tuples
return [(Version(r[0]), pd.DataFrame.from_records(r[1])) for r in records]
def load_payloads(self, record_type=None, *sources, **kwargs):
"""
Reads the contents of MDS payload files.
Parameters:
record_type: str, optional
The type of MDS Provider record.
By default get payloads of each type.
sources: str, Path, list, optional
One or more paths to (directories containing) MDS payload (JSON) files.
Directories are expanded such that all corresponding files within are read.
URLs pointing to JSON files are also supported.
flatten: bool, optional
True (default) to flatten the final result from all sources into a list of dicts.
False to keep each result as-is from the source.
headers: dict, optional
A dict of headers to send with requests made to URL paths.
Could also be a dict mapping an URL path to headers for that path.
ls: callable(sources=list): tuple (files: list, urls: list), optional
A function that receives a list of urllib.parse.ParseResult, and returns
a tuple of a list of valid files, and a list of valid URLs to be read from.
Additional keyword arguments are passed through to json.load().
Raise:
IndexError
When no sources have been specified.
Return:
list
With a single file source, or multiple sources and flatten=True, a list of Provider payload dicts.
With multiple sources and flatten=False, a list of the raw contents of each file.
"""
sources = [self._parse(s) for s in sources]
# record_type is not a schema type, but a data source
if record_type and record_type not in SCHEMA_TYPES:
sources.append(self._parse(record_type))
record_type = None
if len(sources) == 0:
sources.extend(self._sources)
if len(sources) == 0:
raise IndexError("There are no sources to read from.")
record_type = record_type or self.record_type
flatten = kwargs.pop("flatten", True)
headers = kwargs.pop("headers", {})
# obtain a list of file Paths and URL str to read
ls = kwargs.pop("ls", self.ls)
files, urls = ls(sources)
# load from each file/URL pointer into a composite list
data = []
data.extend([json.loads(f.read_text(), **kwargs) for f in files])
data.extend([requests.get(u, headers=headers.get(u, headers)).json() for u in urls])
# filter out payloads with non-matching record_type
if record_type:
filtered = []
data_key = Schema(record_type).data_key
for payload in data:
if isinstance(payload, list):
filtered.extend(filter(lambda p: data_key in p["data"], payload))
elif "data" in payload and data_key in payload["data"]:
filtered.append(payload)
data = filtered
# flatten any sublists
if flatten:
flattened = []
for payload in data:
if isinstance(payload, list):
flattened.extend(payload)
else:
flattened.append(payload)
data = flattened
return data
def load_records(self, record_type=None, *sources, **kwargs):
"""
Reads the contents of MDS payload files into tuples of (Version, list).
Parameters:
record_type: str, optional
The type of MDS Provider record.
sources: str, optional
One or more paths to (directories containing) MDS payload (JSON) files.
flatten: bool, optional
True (default) to flatten the final result from all sources into a single list.
False to keep each result separate.
headers: dict, optional
A dict of headers to send with requests made to URL paths.
Could also be a dict mapping an URL path to headers for that path.
ls: callable(sources=list): list, optional
A function that receives a list of urllib.parse.ParseResult, and returns the
complete list of file Path objects and URL str to be read.
Raise:
UnexpectedVersionError
When flatten=True and a version mismatch is found amongst the data.
ValueError
When neither record_type or instance.record_type is provided.
Return:
tuple (Version, list)
With flatten=True, a (Version, list) tuple.
list
With flatten=False, a list of (Version, list) tuples with length equal to the
total number of payloads across all sources.
"""
record_type = self._record_type_or_raise(record_type)
flatten = kwargs.pop("flatten", True)
# obtain unmodified payloads
kwargs["flatten"] = False
payloads = self.load_payloads(record_type, *sources, **kwargs)
if len(payloads) < 1:
return payloads
# get the version from the initial payload
if isinstance(payloads[0], list):
version = Version(payloads[0][0]["version"])
else:
version = Version(payloads[0]["version"])
# collect versions and data from each payload
data_key = Schema(record_type).data_key
_payloads = []
for payload in payloads:
if not isinstance(payload, list):
payload = [payload]
for page in payload:
data = page["data"][data_key]
# insert last_updated and ttl data from outer payload into each vehicle record
if record_type == VEHICLES:
last_updated, ttl = page["last_updated"], page["ttl"]
for item in data:
item["last_updated"] = last_updated
item["ttl"] = ttl
_payloads.append((page["version"], data))
if flatten:
if not all([Version(v) == version for v,_ in _payloads]):
# find the first non-matching version and raise
unexpected = [Version(v) for v,_ in _payloads if Version(v) != version][0]
raise UnexpectedVersionError(unexpected, version)
# return the version, records tuple
return version, [item for _,data in _payloads for item in data]
else:
# list of version, records tuples
return [(Version(r[0]), r[1]) for r in _payloads]
@classmethod
def _filename(cls, **kwargs):
"""
Generate a filename from the given parameters.
"""
record_type = kwargs.get("record_type", None)
payloads = kwargs.get("payloads", [])
extension = kwargs.get("extension", ".json")
payload = kwargs.get("payload", None)
# is there a single record_type in these payloads that we should use?
record_types = set([list(p["data"].keys())[0] for p in payloads])
if record_type is None and len(record_types) == 1:
record_type = record_types.pop()
# no record_type specified, generate filename from payload hash
if record_type is None:
data = json.dumps(payload or payloads).encode()
shadigest = hashlib.sha256(data).hexdigest()
return f"{shadigest[0:7]}{extension}"
data_key = Schema(record_type).data_key
# find time boundaries from the data
if record_type in [STATUS_CHANGES, EVENTS]:
time_key = "event_time"
elif record_type == TRIPS:
time_key = "end_time"
elif record_type == VEHICLES:
time_key = "last_event_time"
times = [d[time_key] for p in payloads for d in p["data"][data_key]]
if all([isinstance(t, datetime.datetime) for t in times]):
start = min(times)
end = max(times)
else:
decoder = TimestampDecoder()
start = decoder.decode(min(times))
end = decoder.decode(max(times))
# clip to hour of day, offset if they are the same
start = datetime.datetime(start.year, start.month, start.day, start.hour)
end = datetime.datetime(end.year, end.month, end.day, end.hour)
if start == end:
end = end + datetime.timedelta(hours=1)
encoder = TimestampEncoder(date_format="%Y%m%dT%H0000Z")
providers = set([d["provider_name"] for p in payloads for d in p["data"][data_key]])
return f"{'_'.join(providers)}_{record_type}_{encoder.encode(start)}_{encoder.encode(end)}{extension}"
@classmethod
def _ls(cls, sources):
"""
Create a tuple of lists of valid file Paths and URLs from a list of urllib.parse.ParseResult.
"""
# separate into files and directories and urls
files = [pathlib.Path(f.path) for f in sources if cls._isfile(f)]
dirs = [pathlib.Path(d.path) for d in sources if cls._isdir(d)]
urls = [urllib.parse.urlunparse(u) for u in sources if cls._isurl(u)]
# expand into directories
files.extend([f for ls in [d.glob("*.json") for d in dirs] for f in ls])
return files, urls