1 files changed, 45 insertions, 15 deletions
diff --git a/src/ee/ds/__init__.py b/src/ee/ds/__init__.py
index aff5b64..5899d28 100644
--- a/src/ee/ds/__init__.py
+++ b/src/ee/ds/__init__.py
@@ -2,6 +2,7 @@ import configparser
 import csv
 import logging
 import os
+import shutil
 from functools import total_ordering
 from pathlib import Path
 from typing import MutableMapping, Optional, List, Tuple, Union, Iterator
@@ -64,12 +65,14 @@ class Object(object):
     def key(self):
         return self._key
 
-    def set(self, key: str, value: str):
+    def set(self, key: str, value: str) -> "Object":
         if self._ds._frozen:
             raise Exception("This data set is frozen")
         idx = self._ot.index_of(key, create=True)
         self._data.insert(idx, value)
 
+        return self
+
     def _set_from_object(self, other: "Object"):
         for k in other._ot.fields:
             self.set(k, other.get(k))
@@ -89,7 +92,7 @@ class Object(object):
 
 
 class DataSet(object):
-    def __init__(self, name):
+    def __init__(self, name: Optional[str] = None):
         self._name = name
         self._object_types = {}  # type: MutableMapping[str, ObjectType]
         self._objects_by_type = {}  # type: MutableMapping[ObjectType, MutableMapping[str, Object]]
@@ -97,6 +100,8 @@ class DataSet(object):
 
     @property
     def name(self):
+        if not self._name:
+            raise Exception("Unnamed data set")
         return self._name
 
     def __len__(self):
@@ -210,11 +215,11 @@ class DataSetManager(object):
         except KeyError:
             return self._basedir / ds_name / "data-set.ini"
 
-    def create_rw(self, name, inputs: List[str] = None) -> "LazyDataSet":
-        return LazyDataSet(self, False, name, inputs if inputs else [])
+    def create_rw(self, name, clean: bool) -> "LazyRwDataSet":
+        return LazyRwDataSet(self, name, clean)
 
-    def create_ro(self, inputs: List[str]) -> "LazyDataSet":
-        return LazyDataSet(self, True, None, inputs)
+    def create_ro(self, inputs: List[str]) -> "LazyRoDataSet":
+        return LazyRoDataSet(self, inputs)
 
     def add_ds(self, ds_type: str, name: str, object_type: str, path: str = None):
         if ds_type == "csv":
@@ -337,27 +342,52 @@ class DataSetManager(object):
         with open(path, "w") as f:
             ini.write(f)
 
+    def remove(self, name: str):
+        try:
+            object_type, path = self._csv[name]
+            os.remove(str(path))
+        except KeyError:
+            shutil.rmtree(self._basedir / name)
+
 
-class LazyDataSet(object):
-    def __init__(self, dsm: DataSetManager, freeze: bool, name, inputs):
+class LazyRoDataSet(object):
+    def __init__(self, dsm: DataSetManager, inputs):
         self._dsm = dsm
-        self._freeze = freeze
-        self._name = name
         self._inputs = inputs
 
     def __enter__(self) -> DataSet:
         # logger.info("enter: name={}, inputs={}".format(self._name, self._inputs))
-        ds = DataSet(self._name)
+        ds = DataSet()
         for name in self._inputs:
             ds = ds.merge(self._dsm.load(name, freeze=True))
 
-        if self._freeze:
-            ds.freeze()
+        ds.freeze()
+
+        self._ds = ds
+        return ds
 
+    def __exit__(self, *args):
+        return False
+
+
+class LazyRwDataSet(object):
+    def __init__(self, dsm: DataSetManager, name, clean):
+        self._dsm = dsm
+        self._name = name
+        self._clean = clean
+
+    def __enter__(self) -> DataSet:
+        cookie = self._dsm.cookie_for_ds(self._name)
+
+        if cookie.exists():
+            if not self._clean:
+                raise IOError("DataSet already exists: {}, cookie={}".format(self._name, cookie))
+            self._dsm.remove(self._name)
+
+        ds = DataSet(self._name)
         self._ds = ds
         return ds
 
     def __exit__(self, *args):
-        if not self._freeze:
-            self._dsm.store(self._ds)
+        self._dsm.store(self._ds)
         return False