-
Notifications
You must be signed in to change notification settings - Fork 12
/
util.py
165 lines (132 loc) · 5.05 KB
/
util.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
import itertools
import json
from datetime import date
from decimal import Decimal
from functools import wraps
from urllib.parse import parse_qs, urlencode, urlsplit
from ijson import ObjectBuilder, utils
def _pluck_filename(opts):
if opts.package_pointer:
parts = ['pluck', 'package', opts.package_pointer[1:].replace('/', '-')]
else:
parts = ['pluck', 'release', opts.release_pointer[1:].replace('/', '-')]
return f"{'-'.join(parts)}.csv"
def components(start, stop=None):
"""
Returns a function that returns the selected non-empty path components, excluding the ``.json`` extension.
>>> components(-1)('http://example.com/api/planning.json')
'planning'
>>> components(-2, -1)('http://example.com/api/planning/package.json')
'planning'
"""
def wrapper(url):
value = '-'.join(list(filter(None, urlsplit(url).path.split('/')))[start:stop])
if value.endswith('.json'):
return value[:-5]
return value
return wrapper
def parameters(*keys):
"""
Returns a function that returns the selected query string parameters.
>>> parameters('page')('http://example.com/api/packages.json?page=1')
'page-1'
>>> parameters('year', 'page')('http://example.com/api/packages.json?year=2000&page=1')
'year-2000-page-1'
"""
def wrapper(url):
query = parse_qs(urlsplit(url).query)
return '-'.join(s for key in keys for value in query[key] for s in [key, value])
return wrapper
def join(*functions):
"""
Returns a function that joins the given functions' outputs.
>>> join(components(-1), parameters('page'))('http://example.com/api/planning.json?page=1')
'planning-page-1'
"""
def wrapper(url):
return '-'.join(function(url) for function in functions)
return wrapper
def handle_http_error(decorated):
"""
A decorator for spider parse methods.
Yields a :class:`~kingfisher_scrapy.items.FileError` for successful HTTP status codes.
"""
@wraps(decorated)
def wrapper(self, response):
if self.is_http_success(response):
yield from decorated(self, response)
else:
yield self.build_file_error_from_response(response)
return wrapper
# https://stackoverflow.com/questions/34898525/generate-list-of-months-between-interval-in-python
def date_range_by_month(start, stop):
"""
Yields the first day of the month from the ``start`` to the ``stop`` dates, in reverse chronological order.
"""
def number_of_months(d):
return 12 * d.year + d.month
for months in reversed(range(number_of_months(start) - 1, number_of_months(stop))):
year, month = divmod(months, 12)
yield date(year, month + 1, 1)
def date_range_by_year(start, stop):
"""
Returns the year from the ``start`` to the ``stop`` years, in reverse chronological order.
"""
return reversed(range(start, stop + 1))
def replace_parameter(url, key, value):
"""
Returns a URL after updating the query string parameter.
"""
parsed = urlsplit(url)
query = parse_qs(parsed.query)
if value is None:
query.pop(key, None)
else:
query[key] = [value]
return parsed._replace(query=urlencode(query, doseq=True)).geturl()
@utils.coroutine
def items_basecoro(target, prefix, map_type=None, skip_key=None):
"""
This is copied from ``ijson/common.py``. A ``skip_key`` argument is added. If the ``skip_key`` is in the current
path, the current event is skipped. Otherwise, the method is identical.
"""
while True:
current, event, value = (yield)
if skip_key and skip_key in current:
continue
if current == prefix:
if event in ('start_map', 'start_array'):
builder = ObjectBuilder(map_type=map_type)
end_event = event.replace('start', 'end')
while (current, event) != (prefix, end_event):
builder.event(event, value)
current, event, value = (yield)
del builder.containers[:]
target.send(builder.value)
else:
target.send(value)
def items(events, prefix, map_type=None, skip_key=None):
"""
This is copied from ``ijson/common.py``. A ``skip_key`` argument is added, which is passed as a keyword argument to
:meth:`~kingfisher_scrapy.util.items_basecoro`. Otherwise, the method is identical.
"""
return utils.coros2gen(events,
(items_basecoro, (prefix,), {'map_type': map_type, 'skip_key': skip_key}) # noqa: E128
)
def default(obj):
"""
Dumps JSON to a string, and returns it.
"""
if isinstance(obj, Decimal):
return float(obj)
try:
iterable = iter(obj)
except TypeError:
pass
else:
return list(iterable)
return json.JSONEncoder().default(obj)
# See `grouper` recipe: https://docs.python.org/3.8/library/itertools.html#recipes
def grouper(iterable, n, fillvalue=None):
args = [iter(iterable)] * n
return itertools.zip_longest(*args, fillvalue=fillvalue)