Moved dataset loading function to dataset module

3 files changed, 29 insertions, 23 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])]
diff --git a/szilagyi/nomogram.py b/szilagyi/nomogram.py
index 9d76127..91d75c2 100644
--- a/szilagyi/nomogram.py
+++ b/szilagyi/nomogram.py
@@ -1,7 +1,4 @@
-import csv
 import math
-import os
-import re
 from collections import deque
 
 
@@ -15,25 +12,6 @@ class Vector(complex):
 			raise IndexError
 
 
-def load(directory):
-	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 file in os.listdir(directory):
-			match = re.match(r"SWI_(-?\d+)\.csv", file)
-			if match:
-				yield int(match.group(1)), os.path.join(directory, file)
-
-	return [(x, _load(y)) for x, y in sorted(_files(directory), key=lambda x: x[0])]
-
-
 def look_downwards(data, x, start):
 	for i in range(start, 0, -1):
 		if data[i - 1][0] < x:
diff --git a/szilagyi/plots.py b/szilagyi/plots.py
index 8b325d6..9e0c2c5 100644
--- a/szilagyi/plots.py
+++ b/szilagyi/plots.py
@@ -2,6 +2,7 @@ import os
 
 import matplotlib.pyplot as plot
 
+from . import _dataset
 from .nomogram import *
 
 
@@ -14,7 +15,7 @@ def grid(start, end, steps):
 
 
 if __name__ == "__main__":
-	swis = load(os.path.join(os.path.dirname(__file__), "_dataset"))
+	swis = _dataset.load()
 
 	def _swi(x, y):  # dt, depth
 		segments = find_boundary_curves(swis, x, y)