add overrides to *like functions

scikit-hep · Aug 25, 2018 · 58e5035 · 58e5035
1 parent 3dc8336
commit 58e5035
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 20 deletions.
diff --git a/awkward/array/jagged.py b/awkward/array/jagged.py
@@ -187,29 +187,29 @@ def deepcopy(self, starts=None, stops=None, content=None):
         out._parents = awkward.util.deepcopy(out._parents)
         return out
 
-    def like(self, array):
+    def like(self, array, **overrides):
         if isinstance(self._content, awkward.util.numpy.ndarray):
             return self.copy(content=array)
         else:
-            return self.copy(content=self._content.like(array))
+            return self.copy(content=self._content.like(array, **overrides))
 
-    def empty_like(self):
+    def empty_like(self, **overrides):
         if isinstance(self._content, awkward.util.numpy.ndarray):
             return self.copy(content=awkward.util.numpy.empty_like(self._content))
         else:
-            return self.copy(content=self._content.empty_like())
+            return self.copy(content=self._content.empty_like(**overrides))
 
-    def zeros_like(self):
+    def zeros_like(self, **overrides):
         if isinstance(self._content, awkward.util.numpy.ndarray):
             return self.copy(content=awkward.util.numpy.zeros_like(self._content))
         else:
-            return self.copy(content=self._content.zeros_like())
+            return self.copy(content=self._content.zeros_like(**overrides))
 
-    def ones_like(self):
+    def ones_like(self, **overrides):
         if isinstance(self._content, awkward.util.numpy.ndarray):
             return self.copy(content=awkward.util.numpy.ones_like(self._content))
         else:
-            return self.copy(content=self._content.ones_like())
+            return self.copy(content=self._content.ones_like(**overrides))
 
     @property
     def starts(self):

diff --git a/awkward/array/objects.py b/awkward/array/objects.py
@@ -61,29 +61,45 @@ def deepcopy(self, content=None, generator=None, args=None, kwargs=None):
         out._content = awkward.util.deepcopy(out._content)
         return out
 
-    def like(self, array):
+    def like(self, array, **overrides):
+        mine = {}
+        mine["generator"] = overrides.pop("generator", self._generator)
+        mine["args"] = overrides.pop("args", self._args)
+        mine["kwargs"] = overrides.pop("kwargs", self._kwargs)
         if isinstance(self._content, awkward.util.numpy.ndarray):
             return self.copy(content=array)
         else:
-            return self.copy(content=self._content.like(array))
+            return self.copy(content=self._content.like(array, **overrides))
 
-    def empty_like(self):
+    def empty_like(self, **overrides):
+        mine = {}
+        mine["generator"] = overrides.pop("generator", self._generator)
+        mine["args"] = overrides.pop("args", self._args)
+        mine["kwargs"] = overrides.pop("kwargs", self._kwargs)
         if isinstance(self._content, awkward.util.numpy.ndarray):
-            return self.copy(content=awkward.util.numpy.empty_like(self._content))
+            return self.copy(content=awkward.util.numpy.empty_like(self._content), **mine)
         else:
-            return self.copy(content=self._content.empty_like())
+            return self.copy(content=self._content.empty_like(**overrides), **mine)
 
-    def zeros_like(self):
+    def zeros_like(self, **overrides):
+        mine = {}
+        mine["generator"] = overrides.pop("generator", self._generator)
+        mine["args"] = overrides.pop("args", self._args)
+        mine["kwargs"] = overrides.pop("kwargs", self._kwargs)
         if isinstance(self._content, awkward.util.numpy.ndarray):
-            return self.copy(content=awkward.util.numpy.zeros_like(self._content))
+            return self.copy(content=awkward.util.numpy.zeros_like(self._content), **mine)
         else:
-            return self.copy(content=self._content.zeros_like())
+            return self.copy(content=self._content.zeros_like(**overrides), **mine)
 
-    def ones_like(self):
+    def ones_like(self, **overrides):
+        mine = {}
+        mine["generator"] = overrides.pop("generator", self._generator)
+        mine["args"] = overrides.pop("args", self._args)
+        mine["kwargs"] = overrides.pop("kwargs", self._kwargs)
         if isinstance(self._content, awkward.util.numpy.ndarray):
-            return self.copy(content=awkward.util.numpy.ones_like(self._content))
+            return self.copy(content=awkward.util.numpy.ones_like(self._content), **mine)
         else:
-            return self.copy(content=self._content.ones_like())
+            return self.copy(content=self._content.ones_like(**overrides), **mine)
 
     @property
     def content(self):

diff --git a/awkward/array/table.py b/awkward/array/table.py
@@ -192,7 +192,7 @@ def deepcopy(self, content=None):
             out._base = None
         return out
 
-    def empty_like(self):
+    def empty_like(self, **overrides):
         out = self.__class__.__new__(self.__class__)
         out._view = None
         out._base = None