diff options
| -rw-r--r-- | szilagyi/_dataset/__init__.py | 18 |
1 files changed, 10 insertions, 8 deletions
diff --git a/szilagyi/_dataset/__init__.py b/szilagyi/_dataset/__init__.py index 8ff88ae..e91251e 100644 --- a/szilagyi/_dataset/__init__.py +++ b/szilagyi/_dataset/__init__.py @@ -7,21 +7,20 @@ class Vector(complex): def __getitem__(self, index): if index == 0: return self.real - elif index == 1: + if index == 1: return self.imag - else: - raise IndexError + raise IndexError def _load_from(root): - def _read(iterable): + def _vectors(iterable): for x, y in iterable: yield Vector(float(x), float(y)) def _load(filename): - with open(filename) as fd: - reader = csv.reader(fd) - return list(_read(reader)) + with open(filename, encoding='utf-8') as fp: + reader = csv.reader(fp) + return list(_vectors(reader)) def _files(directory): for filename in os.listdir(directory): @@ -29,7 +28,10 @@ def _load_from(root): if match: yield int(match.group(1)), os.path.join(directory, filename) - return [(x, _load(y)) for x, y in sorted(_files(root), key=lambda x: x[0])] + def _first(pair): + return pair[0] + + return [(index, _load(path)) for index, path in sorted(_files(root), key=_first)] INDICES = _load_from(os.path.dirname(os.path.abspath(__file__))) |