-
-
Notifications
You must be signed in to change notification settings - Fork 27
Expand file tree
/
Copy pathDataFileInterface.py
More file actions
executable file
·292 lines (229 loc) · 10.3 KB
/
Copy pathDataFileInterface.py
File metadata and controls
executable file
·292 lines (229 loc) · 10.3 KB
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
from pathlib import Path
from typing import Any, Iterable, Iterator
from yaml import safe_load, dump
from modules import global_objects
from modules.Debug import log
from modules.EpisodeInfo import EpisodeInfo
from modules.SeriesInfo import SeriesInfo
from modules.Title import Title
class DataFileInterface:
"""
This class is used to interface with a show's data file. And can be
used for reading from and writing to the files for the purpose of
adding new or reading existing episode data.
"""
"""Default name for a data file of episode information"""
GENERIC_DATA_FILE_NAME = 'data.yml'
__slots__ = ('series_info', 'file', '_ignore_preferred_titles')
def __init__(self,
series_info: SeriesInfo,
data_file: Path,
*,
ignore_preferred_titles: bool = False,
) -> None:
"""
Constructs a new instance of the interface for the specified
data file. This also creates the parent directories for the data
file if they do not exist.
Args:
series_info: Series associated with this interface. For
logging only.
data_file: Path to the data file to interface with.
ignore_preferred_titles: Whether to ignore preferred titles
when reading this data file.
"""
# Store the SeriesInfo and data file
self.series_info = series_info
self.file = data_file
self._ignore_preferred_titles = ignore_preferred_titles
# Create parent directories if necessary
if not self.file.exists():
data_file.parent.mkdir(parents=True, exist_ok=True)
def __repr__(self) -> str:
"""Returns an unambiguous string representation of the object."""
return (f'<DataFileInterface series_info={self.series_info}, '
f'file={self.file.resolve()}>')
def __read_data(self) -> dict[str, dict[float, dict]]:
"""
Read this interface's data from file. Returns an empty
dictionary if the file does not exist, is misformatted, or if
'data' key is missing.
Returns:
Contents under 'data' key of this interface's file.
"""
# If the file DNE, return empty dictionary
if not self.file.exists():
return {}
# Read file
with self.file.open('r', encoding='utf-8') as file_handle:
try:
yaml = safe_load(file_handle)
except Exception as e: # pylint: disable=broad-except
log.error(f'Error reading datafile:\n{e}\n')
return {}
# If the top-level key is not 'data', error and return empty dictionary
if 'data' not in yaml:
log.error(f'Datafile "{self.file.resolve()}" missing "data" key')
return {}
if not isinstance(yaml['data'], dict):
log.error(f'Data in "{self.file.resolve()}" is invalid')
return {}
return yaml['data']
def __write_data(self, yaml: dict[str, Any]) -> None:
"""
Write the given YAML data to this interface's file. This puts
all data under the 'data' key.
Args:
yaml: YAML dictionary to write to file.
"""
# Write updated data with this entry added
with self.file.open('w', encoding='utf-8') as file_handle:
dump({'data': yaml}, file_handle, allow_unicode=True, width=100)
def read(self) -> Iterator[tuple[dict[str, Any], set[str]]]:
"""
Read the data file for this object, yielding each valid row.
Returns:
Yields a dictionary for each entry in this datafile. The
dictionary has a key 'episode_info' with an EpisodeInfo
object, and arbitrary keys for all other data found within
the entry's YAML.
"""
# Read yaml, returns {} if empty/DNE
yaml = self.__read_data()
# Iterate through each season
for season, season_data in yaml.items():
# Skip season if number cannot be parsed
try:
season_number = int(season.rsplit(' ', 1)[-1])
except (IndexError, ValueError):
log.error(f'Season {season} of the {self.series_info} datafile '
f'is invalid - must be like "Season 1"')
continue
# Iterate through each episode of this season
for episode_number, episode_data in season_data.items():
# Ensure indices are integers
abs_number = episode_data.get('abs_number')
try:
episode_number = int(episode_number)
if abs_number is not None:
abs_number = int(abs_number)
except ValueError:
log.exception(f'Episode data for S{season_number:02} is '
f'invalid - episode and absolute number must '
f'be integers')
continue
# If title is missing (or no subkeys at all..) error
if (not isinstance(episode_data, dict)
or ('title' not in episode_data and
'preferred_title' not in episode_data)):
log.error(f'S{season_number:02}E{episode_number:02} of the '
f'{self.series_info} datafile is missing a title')
continue
# Get existing keys for this episode
given_keys = set(episode_data)
# If translated title is available, prefer that
original_title = episode_data.pop('title', None)
if self._ignore_preferred_titles:
title = original_title
else:
title = episode_data.get('preferred_title', original_title)
# Ensure Title can be created
try:
title_obj = Title(title, original_title=original_title)
except TypeError:
log.exception(
f'Title for S{season_number:02}E{episode_number:02} of '
'the {self.series_info} datafile is invalid'
)
continue
# Construct EpisodeInfo object for this entry
episode_info = global_objects.info_set.get_episode_info(
self.series_info,
title_obj,
season_number,
episode_number,
abs_number,
emby_id=episode_data.pop('emby_id', None),
imdb_id=episode_data.pop('imdb_id', None),
tmdb_id=episode_data.pop('tmdb_id', None),
tvdb_id=episode_data.pop('tvdb_id', None),
tvrage_id=episode_data.pop('tvrage_id', None),
)
# Add any additional, unexpected keys from the YAML
data = {'episode_info': episode_info}
data.update(episode_data)
yield data, given_keys
def __info_as_entry(self, episode_info: EpisodeInfo) -> dict[str, Any]:
"""
Get the given EpisodeInfo object as it's equivalent YAML entry.
Args:
episode_info: EpisodeInfo to get the entry of.
Returns:
Dictionary to write under episode number key of the given
info. Possible keys are 'title', and 'abs_number'.
"""
entry = {'title': episode_info.title.title_yaml}
if episode_info.abs_number is not None:
entry['abs_number'] = episode_info.abs_number
return entry
def add_data_to_entry(self,
episode_info: EpisodeInfo,
**new_data: dict[str, Any],
) -> None:
"""
Add any generic data to the YAML entry associated with this
EpisodeInfo.
Args:
episode_info: Episode Info to add to YAML.
new_data: Generic new data to write.
"""
yaml = self.__read_data()
# Verify this entry already exists, warn and exit if not
season_key = f'Season {episode_info.season_number}'
if (season_key not in yaml
or episode_info.episode_number not in yaml[season_key]):
log.error(f'Cannot add data to entry for {episode_info} in '
f'"{self.file.resolve()}" - entry does not exist')
return None
# Add new data
yaml[season_key][episode_info.episode_number].update(new_data)
# Write updated data
self.__write_data(yaml)
return None
def add_many_entries(self, new_episodes: Iterable[EpisodeInfo]) -> None:
"""
Adds many entries at once. This only reads and writes from this
interface's file once.
Args:
new_episodes: Iterable of EpisodeInfo objects to write.
"""
# If no new episodes are being added, exit
if len(new_episodes) == 0:
return None
# Read yaml
yaml = self.__read_data()
# Go through each episode to possibly add to file
added = {'count': 0, 'info': None}
for episode_info in new_episodes:
# Create blank season data if this key doesn't exist
season_key = f'Season {episode_info.season_number}'
if season_key not in yaml:
yaml[season_key] = {}
# Construct episode data
data = self.__info_as_entry(episode_info)
# Add episode data to existing entry or create new entry for episode
added = {'count': added['count'] + 1, 'info': episode_info}
if yaml[season_key].get(episode_info.episode_number) is not None:
yaml[season_key][episode_info.episode_number].update(data)
else:
yaml[season_key][episode_info.episode_number] = data
# If nothing was added, exit - otherwise log to user
if (count := added['count']) == 0:
return None
if count > 1:
log.info(f'Added {count} episodes to "{self.file.parent.name}"')
else:
log.info(f'Added {added["info"]} to "{self.file.parent.name}"')
# Write updated yaml
self.__write_data(yaml)
return None