Moved dataset loading function to dataset module

1 files changed, 27 insertions, 0 deletions

diff --git a/szilagyi/_dataset/__init__.py b/szilagyi/_dataset/__init__.py
new file mode 100644
index 0000000..cfa1aee
--- /dev/null
+++ b/szilagyi/_dataset/__init__.py
@@ -0,0 +1,27 @@
+import csv
+import os
+import re
+
+from ..nomogram import Vector
+
+
+ROOT = os.path.dirname(os.path.abspath(__file__))
+
+
+def load():
+	def _read(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))
+
+	def _files(directory):
+		for filename in os.listdir(directory):
+			match = re.match(r"SWI_(-?\d+)\.csv", filename)
+			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])]