diff --git a/fluents/dataset.py b/fluents/dataset.py index a827321..8722c6e 100644 --- a/fluents/dataset.py +++ b/fluents/dataset.py @@ -444,6 +444,11 @@ def write_ftsv(fd, ds): @param ds: The dataset to be written. The function handles datasets of these classes: Dataset, CategoryDataset and GraphDataset """ + opened = False + if isinstance(fd, str): + fd = open(fd, 'w') + opened = True + # Write header information if isinstance(ds, CategoryDataset): type = 'category' @@ -475,6 +480,8 @@ def write_ftsv(fd, ds): print >> fd, "%s\t" % m[j, i], print >> fd + if opened: + fd.close() def read_ftsv(fd): """Read a dataset in fluents tab separated values (ftsv) form and return it. @@ -483,6 +490,11 @@ def read_ftsv(fd): @return: A Dataset, CategoryDataset or GraphDataset depending on the information read. """ + opened = False + if isinstance(fd, str): + fd = open(fd) + opened = True + split_re = re.compile('^#\s*(\w+)\s*:\s*(.+)') dimensions = [] identifiers = {} @@ -548,5 +560,9 @@ def read_ftsv(fd): else: ds = Dataset(matrix, dims, name) + if opened: + fd.close() + return ds +