Merge pull request #2 from DolphDev/latest

Push Fixes and Additions to master

psv/__init__.py

@@ -1,5 +1,6 @@
 from .core.objects.apiobjects import MainSelection
-from .core.objects import BaseRow
+from .core.objects import BaseRow, banned_columns
+from .core.exceptions.messages import LoadingMsg as msg
 
 
 import csv
@@ -17,7 +18,7 @@ def csv_size_limit(size):
 
 
 def load(f, cls=BaseRow, outputfile=None, delimiter=",", quotechar='"', mode='r', buffering=-1,
-         encoding="utf-8", errors=None, newline=None, closefd=True, opener=None, typetranfer=True,
+         encoding="utf-8", errors=None, newline=None, closefd=True, opener=None, typetransfer=True,
          csv_size_max=None, csv_max_row=None):
     """Loads a file into psv
 
@@ -35,66 +36,82 @@ def load(f, cls=BaseRow, outputfile=None, delimiter=",", quotechar='"', mode='r'
                 csvfile, delimiter=delimiter, quotechar=quotechar)
             api = MainSelection(data, columns=column_names(csvfile.name, cls, quotechar, delimiter,
                                                            mode, buffering, encoding, errors, newline, closefd, opener),
-                                outputfiled=outputfile, cls=cls, typetranfer=typetranfer)
+                                outputfiled=outputfile, cls=cls, typetransfer=typetransfer)
     else:
         with f if isinstance(f, io._io._IOBase) else open(f, mode=mode, buffering=buffering,
                                                           encoding=encoding, errors=errors, newline=newline, closefd=closefd, opener=opener) as csvfile:
             data = itertools.islice(csv.DictReader(
                 csvfile, delimiter=delimiter, quotechar=quotechar), csv_max_row)
             api = MainSelection(data, columns=column_names(csvfile.name, cls, quotechar, delimiter,
                                                            mode, buffering, encoding, errors, newline, closefd, opener),
-                                outputfiled=outputfile, cls=cls, typetranfer=typetranfer)
+                                outputfiled=outputfile, cls=cls, typetransfer=typetransfer)
     return api
 
 
 def loaddir(f, cls=BaseRow, outputfile=None, delimiter=",", quotechar='"', mode='r', buffering=-1,
-            encoding="utf-8", errors=None, newline=None, closefd=True, opener=None, typetranfer=True,
-            csv_size_max=None):
+            encoding="utf-8", errors=None, newline=None, closefd=True, opener=None, typetransfer=True,
+            csv_size_max=None, filetype="*.csv"):
     """Loads a directory of .csv files"""
     if csv_size_max:
         csv_size_limit(csv_size_max)
     data = []
     columns = None
-    for files in glob.glob(f+"*.csv"):
+    for files in glob.glob(f+filetype):
         if not columns:
             columns = column_names(files)
         with open(files, mode=mode, buffering=buffering,
                   encoding=encoding, errors=errors, newline=newline, closefd=closefd, opener=opener) as csvfile:
             data = data + list(csv.DictReader(csvfile,
                                               delimiter=delimiter, quotechar=quotechar))
-    return MainSelection(data, columns=columns, outputfiled=outputfile, cls=cls, typetranfer=typetranfer)
+    forbidden_columns(columns)
+    return MainSelection(data, columns=columns, outputfiled=outputfile, cls=cls, typetransfer=typetransfer)
 
 
 def loads(csvdoc, columns=None, cls=BaseRow, outputfile=None, delimiter=",", quotechar='"',
-          typetranfer=True, csv_size_max=None, newline="\n"):
+          typetransfer=True, csv_size_max=None, newline="\n"):
+    was_str = False
     if csv_size_max:
         csv_size_limit(csv_size_max)
     if isinstance(csvdoc, str):
+        was_str = True
         data = csv.DictReader(csvdoc.split(newline),
                               delimiter=delimiter, quotechar=quotechar)
         if not columns:
             columns = tuple(next(csv.reader(csvdoc.split(
                 newline), delimiter=delimiter, quotechar=quotechar)))
     else:
         data = csvdoc
+    if columns:
+        forbidden_columns(columns)
+    elif (not columns) and isinstance(csvdoc, dict):
+        forbidden_columns(csvdoc.keys())
     api = MainSelection(data, columns=(
-        columns), outputfiled=outputfile, cls=cls, typetranfer=typetranfer)
+        columns), outputfiled=outputfile, cls=cls, typetransfer=typetransfer)
     return api
 
 
 def new(cls=BaseRow, columns=None, outputfile=None,
         csv_size_max=None):
     if csv_size_max:
         csv_size_limit(csv_size_max)
+    if columns:
+        forbidden_columns(columns)
     return MainSelection(columns=columns, outputfiled=outputfile, cls=cls)
 
 
 def column_names(f, cls=BaseRow, quotechar='"', delimiter=",", mode='r', buffering=-1, encoding="utf-8",
                  errors=None, newline=None, closefd=True, opener=None,
-                 csv_size_max=None):
+                 csv_size_max=None, check_columns=True):
     if csv_size_max:
         csv_size_limit(csv_size_max)
     with open(f, mode=mode, buffering=buffering,
               encoding=encoding, errors=errors, newline=newline, closefd=closefd, opener=opener) as csvfile:
         columns = next(csv.reader(csvfile, delimiter=',', quotechar=quotechar))
+    if check_columns:
+        forbidden_columns(columns)
     return tuple(columns)
+
+def forbidden_columns(columns):
+    for x in columns:
+        if x in banned_columns:
+            raise ValueError(msg.forbidden_column(x))

psv/core/exceptions/messages.py

@@ -5,7 +5,8 @@ class RowObjectMsg:
     flagmessage = "Missing Flag"
     attribute_readonly = "'{classname}' object attribute '{attr}' is read-only. Note: '{attr}' is likely method/class attribute and not part of this row's data"
     attribute_missing = "'{}' has no attribute '{}'"
-    inherited_rows = "Only inherited Rows can use add_valid_attribute"
+    inherited_rows = "Only inherited Objects can use add_valid_attribute"
+    non_valid = "'{}' cannot be shorten to a python valid attribute, and therefore is an invalid column name"
 
 class ApiObjectMsg:
 
@@ -14,3 +15,6 @@ class ApiObjectMsg:
     outputmsg = "A ordered list of columns must be supplied to output the file"
     badgrab = "Empty Grab. Lack of arguments were supplied to grab method."
 
+class LoadingMsg:
+
+    forbidden_column = "'{}' is a forbidden column name, due to it being reserved by psv internally"

psv/core/objects/__init__.py

@@ -1,2 +1,2 @@
-from .rowobjects import BaseRow, cleanup_name
+from .rowobjects import BaseRow, cleanup_name, banned_columns
 from .selections import Selection

psv/core/objects/apiobjects.py

@@ -1,4 +1,4 @@
-from ..objects import BaseRow, Selection
+from ..objects import BaseRow, Selection, cleanup_name
 from ..parsing import parser, parser_addrow
 from ..utils import multiple_index, _index_function_gen
 from ..utils import generate_func
@@ -14,19 +14,34 @@ class MainSelection(Selection):
                  "__columns__", "__columnsmap__",
                  "__rows__", "__apimother__"]
 
-    def __init__(self, csvdict=None, columns=None, cls=BaseRow, parsing=parser, outputfile=None, typetranfer=True, *args, **kwargs):
+    def __init__(self, csvdict=None, columns=None, 
+                 cls=BaseRow, parsing=parser, outputfile=None,
+                 typetranfer=True, *args, **kwargs):
         # Since I close the file after this, the row must be placed into memory
+        rebuild_column_map = False
         self.__apimother__ = self
         self.__outputname__ = outputfile
         self.__columns__ = columns
+        self.__columnsmap__ = {}
+
+        if columns:
+            self.__columnsmap__.update(
+                column_crunch_repeat(self.__columns__))
+        else:
+            rebuild_column_map = True
 
         if csvdict is None:
             csvdict = {}
         if csvdict:
             self.__rows__ = list(
-                parser(csvdict, cls, typetranfer, *args, **kwargs))
+                parser(csvdict, cls, self.__columnsmap__, typetranfer, *args, **kwargs))
+            if rebuild_column_map:
+                self.__columnsmap__.update(
+                    column_crunch_repeat(self.__rows__[0].keys()))
         else:
             self.__rows__ = list()
+        if not self.__columns__:
+            self.__columns__ = tuple()
 
     @property
     def rows(self):
@@ -46,12 +61,73 @@ def columns(self, v):
 
 
     def addrow(self, columns=None, cls=BaseRow, **kwargs):
-        r = parser_addrow(columns if columns else self.__columns__, cls)
+        r = parser_addrow(columns if columns else self.__columns__, cls, self.__columnsmap__)
         self.__rows__.append(r)
         if kwargs:
             for k, v in kwargs.items():
                 r.setcolumn(k, v)
         return r
 
+
+    def addcolumn(self, columnname, columndata="",
+                  add_to_columns=True, clear=True):
+        """Adds a column
+        :param columnname: Name of the column to add.
+        :param columndata: The default value of the new column.
+        :param add_to_columns: Determines whether this column should
+            be added to the internal tracker.
+        :type columnname: :class:`str`
+        :type add_to_columns: :class:`bool`
+        Note: Rows that are being accessed by another thread will error out if clear is True
+        Note:
+        """
+        if columnname in self.columns:
+            raise ValueError(
+                "'{}' column already exists"
+                .format(columnname))
+        for row in self.rows:
+            row.addcolumn(columnname, columndata)
+        if add_to_columns:
+            self.columns += (columnname,)
+            self.__columnsmap__.clear()
+            self.__columnsmap__.update(column_crunch_repeat(self.__columns__))
+        return self
+
     def __delitem__(self, v):
         del self.__rows__[v]
+
+def _column_repeat_dict(columns, clean_ups):
+    master = {}
+    for x in clean_ups:
+        master[x] = None
+    for x in columns:
+        cl_name = cleanup_name(x)
+        if master[cl_name] is None:
+            master[cl_name] = (x,)
+        else:
+            master[cl_name] = master[cl_name] + (x,)
+    return master
+
+def column_crunch_repeat(columns):
+    rv = {}
+    clean_ups = set(cleanup_name(x) for x in columns)
+    ref = _column_repeat_dict(columns, clean_ups)
+    if len(clean_ups) == len(columns):
+        for x in clean_ups:
+            for column in ref[x]:
+                rv.update({x:column})
+    for x in clean_ups:
+        if len(ref[x]) > 1:
+            counter = 0
+            for column in ref[x]:
+                cln = x+"_"+str(counter)
+                while cln in clean_ups:
+                    counter += 1
+                    cln = x+"_"+str(counter)
+                rv.update({cln:column})
+                counter += 1
+        else:
+            for column in ref[x]:
+                rv.update({x:column})
+
+    return rv