Merge pull request #159 from scikit-hep/pad-axis-parameter

add an axis parameter to the pad method

awkward/array/base.py

@@ -473,9 +473,9 @@ def _util_flatten(cls, array, axis):
             return array.reshape(array.shape[:axis] + (-1,) + array.shape[axis + 2:])
 
     @classmethod
-    def _util_pad(cls, array, length, maskedwhen, clip):
+    def _util_pad(cls, array, length, maskedwhen, clip, axis):
         if isinstance(array, AwkwardArray):
-            return array.pad(length, maskedwhen=maskedwhen, clip=clip)
+            return array.pad(length, maskedwhen=maskedwhen, clip=clip, axis=axis)
 
         elif len(array.shape) == 1:
             raise ValueError("pad cannot be applied to scalars")

awkward/array/chunked.py

@@ -630,8 +630,8 @@ def flatten(self, axis=0):
         out.knowchunksizes()
         return out
 
-    def pad(self, length, maskedwhen=True, clip=False):
-        return self.copy(chunks=[self._util_pad(x, length, maskedwhen, clip) for x in self._chunks], chunksizes=self.chunksizes)
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
+        return self.copy(chunks=[self._util_pad(x, length, maskedwhen, clip, axis) for x in self._chunks], chunksizes=self.chunksizes)
 
     def regular(self):
         self._valid()

awkward/array/indexed.py

@@ -247,8 +247,8 @@ def flatten(self, axis=0):
         self._valid()
         return self._util_flatten(self._content[self._index], axis)
 
-    def pad(self, length, maskedwhen=True, clip=False):
-        return self._util_pad(self._content[self._index], length, maskedwhen, clip)
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
+        return self._util_pad(self._content[self._index], length, maskedwhen, clip, axis)
 
     def regular(self):
         self._valid()
@@ -725,8 +725,8 @@ def flatten(self, axis=0):
             out[mask] = content[self._inverse[mask]]
         return out
 
-    def pad(self, length, maskedwhen=True, clip=False):
-        return self._util_pad(self._content.dense, length, maskedwhen, clip)
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
+        return self._util_pad(self._content.dense, length, maskedwhen, clip, axis)
 
     def regular(self):
         self._valid()

awkward/array/jagged.py

@@ -1781,7 +1781,13 @@ def ready(x):
                 table = first.Table.named("tuple", columns1, *columns2)
             return first.JaggedArray(first._starts, first._stops, table)
 
-    def pad(self, length, maskedwhen=True, clip=False):
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
+        if not self._util_isinteger(axis) or axis < 0:
+            raise TypeError("axis must be a non-negative integer (can't count from the end)")
+
+        if axis > 0:
+            return type(self).fromcounts(self.counts, self._util_pad(self._content, length, maskedwhen, clip, axis - 1))
+
         flatstarts = self._starts.reshape(-1)
 
         if clip:

awkward/array/masked.py

@@ -291,8 +291,8 @@ def flattentuple(self):
     def flatten(self, axis=0):
         return self._util_flatten(self._content[self.boolmask(maskedwhen=False)], axis)
 
-    def pad(self, length, maskedwhen=True, clip=False):
-        return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip))
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
+        return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip, axis))
 
     def regular(self):
         self._valid()

awkward/array/objects.py

@@ -262,8 +262,8 @@ def flattentuple(self):
     def flatten(self, axis=0):
         return self.copy(content=self._util_flatten(self._content, axis))
 
-    def pad(self, length, maskedwhen=True, clip=False):
-        return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip))
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
+        return self.copy(content=self._util_pad(self._content, length, maskedwhen, clip, axis))
 
     def regular(self):
         return self.numpy.array(self)
@@ -652,7 +652,13 @@ def flatten(self, axis=0):
         else:
             return self.fromjagged(self.JaggedArray.fromcounts([len(content)], content))
 
-    def pad(self, length, maskedwhen=None, clip=False):
+    def pad(self, length, maskedwhen=None, clip=False, axis=0):
+        if not self._util_isinteger(axis) or axis < 0:
+            raise TypeError("axis must be a non-negative integer (can't count from the end)")
+
+        if axis > 0:
+            raise ValueError("axis too deep for StringArray")
+
         if maskedwhen is None:
             maskedwhen = ord(b" ")
         elif not isinstance(maskedwhen, bytes) or not len(maskedwhen) == 1:

awkward/array/table.py

@@ -767,10 +767,10 @@ def argcross(self, other, nested=False):
     def flatten(self, axis=0):
         raise ValueError("cannot flatten through a Table")
 
-    def pad(self, length, maskedwhen=True, clip=False):
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
         out = self.copy(contents={})
         for n, x in self._contents.items():
-            out[n] = self._util_pad(x, length, maskedwhen, clip)
+            out[n] = self._util_pad(x, length, maskedwhen, clip, axis)
         return out
 
     def regular(self):

awkward/array/union.py

@@ -468,9 +468,9 @@ def flattentuple(self):
     def flatten(self, axis=0):
         raise NotImplementedError("flatten not yet implemented for UnionArray")
 
-    def pad(self, length, maskedwhen=True, clip=False):
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
         self._valid()
-        return self.copy(contents=[self._util_pad(x, length, maskedwhen, clip) for x in self._contents])
+        return self.copy(contents=[self._util_pad(x, length, maskedwhen, clip, axis) for x in self._contents])
 
     def regular(self):
         self._valid()

awkward/array/virtual.py

@@ -439,8 +439,8 @@ def flattentuple(self):
     def flatten(self, axis=0):
         return self._util_flatten(self.array, axis)
 
-    def pad(self, length, maskedwhen=True, clip=False):
-        return self._util_pad(self.array, length, maskedwhen, clip)
+    def pad(self, length, maskedwhen=True, clip=False, axis=0):
+        return self._util_pad(self.array, length, maskedwhen, clip, axis)
 
     def regular(self):
         return self._util_regular(self.array)

awkward/version.py

@@ -4,7 +4,7 @@
 
 import re
 
-__version__ = "0.12.1"
+__version__ = "0.12.2"
 version = __version__
 version_info = tuple(re.split(r"[-\.]", __version__))