Repository: alexhsamuel/ntab

Issue:
assign constant to array
As a convenience, allow,

```py
tbl.a.foo = 42
```

for either an existing or new array.


Comments:

Patch:
diff --git a/ntab/tab.py b/ntab/tab.py
index 3d90a41..d259d13 100644
--- a/ntab/tab.py
+++ b/ntab/tab.py
@@ -9,13 +9,48 @@ from   __future__ import absolute_import, division, print_function, unicode_lite
 from   builtins import *
 import collections
 from   collections import OrderedDict as odict
+import itertools
 import numpy as np
+from   past.builtins import basestring
 import six
 import sys
 
 from   . import fmt, nplib
 from   .lib import *
 
+#-------------------------------------------------------------------------------
+
+def _ensure_array(obj, length):
+    """
+    Ensures `obj` is an ndarray of shape `(length, )`, converting if necessary.
+    """
+    arr = None
+
+    if isinstance(obj, np.ndarray):
+        arr = obj
+
+    if arr is None and not isinstance(obj, basestring):
+        # Convert sequences to arrays.
+        try:
+            len(obj)
+        except:
+            pass
+        else:
+            arr = np.array(obj)
+            
+    # Convert scalars to full arrays.
+    if arr is None:
+        # FIXME: Newer numpy doesn't require explicit dtype
+        dtype = np.array(obj).dtype
+        arr = np.full(length, obj, dtype)
+    
+    if len(arr.shape) != 1:
+        raise ValueError("not one-dimensional")
+    if length is not None and arr.shape != (length, ):
+        raise ValueError("wrong length")
+    return arr
+
+
 #-------------------------------------------------------------------------------
 
 class ArraysObjectProxy(object):
@@ -269,9 +304,9 @@ class Table(object):
             (n, a[sel]) for n, a in self.__arrs.items() )
 
 
-    def __construct(self, arrs):
+    def __construct(self, length, arrs):
+        self.__length = length
         self.__arrs = arrs
-        self.__length = None if len(arrs) == 0 else len(a_value(arrs))
         # Proxies.
         # FIXME: Create lazily?
         self.a          = ArraysObjectProxy(self)
@@ -308,10 +343,26 @@ class Table(object):
         be one-dimensional and the same length.
         """
         arrs = odict(*args, **kw_args)
+
+        # Get the length.
+        length = None
+        for arr in six.itervalues(arrs):
+            try:
+                length = len(arr)
+            except TypeError:
+                pass
+            else:
+                break
+        if length is None and len(arrs) > 0:
+            raise ValueError("no arrs have length")
+
         # Make sure the arrays are all arrays.
-        arrs = odict( (str(n), np.array(a)) for n, a in six.iteritems(arrs) )
+        arrs = odict(
+            (str(n), _ensure_array(a, length)) 
+            for n, a in six.iteritems(arrs) 
+        )
 
-        self.__construct(arrs)
+        self.__construct(length, arrs)
         self.__check(self.__arrs)
 
 
@@ -332,6 +383,7 @@ class Table(object):
         # Construct an instance without calling __init__().
         self = object.__new__(class_)
 
+        length = None if len(arrs) == 0 else len(a_value(arrs))
         self.__construct(arrs)
         if check:
             self.__check(self.__arrs)
@@ -377,11 +429,16 @@ class Table(object):
 
     #---------------------------------------------------------------------------
     # Mutators
-    # FIXME: Make immutable?
 
     def add(self, *args, **kw_args):
+        """
+        Adds or replaces a column.
+        """
         arrs = odict(*args, **kw_args)
-        arrs = odict( (str(n), np.array(a)) for n, a in six.iteritems(arrs) )
+        arrs = odict( 
+            (str(n), _ensure_array(a, self.__length)) 
+            for n, a in six.iteritems(arrs) 
+        )
 
         if len(arrs) == 0:
             # Nothing to do.


Test patch:
diff --git a/ntab/test/test_basic.py b/ntab/test/test_basic.py
index 44e7468..6259508 100644
--- a/ntab/test/test_basic.py
+++ b/ntab/test/test_basic.py
@@ -81,3 +81,18 @@ def test_empty_arrs():
     assert tab.num_rows == 0
 
 
+def test_tab_create_scalar():
+    tab = Table(i=2, x=[3, 4, 5], l="foo")
+    assert tab.num_rows == 3
+    assert list(tab.a.i) == [2, 2, 2]
+    assert list(tab.a.l) == ["foo", "foo", "foo"]
+
+
+def test_tab_add_col_scalar():
+    tab = Table(x=[3, 4, 5])
+    tab.a.i = 2
+    tab.a.l = "foo"
+    assert list(tab.a.i) == [2, 2, 2]
+    assert list(tab.a.l) == ["foo", "foo", "foo"]
+
+