diff --git a/distarray/dist/decorators.py b/distarray/dist/decorators.py index 32accfc9..62b7aaa5 100644 --- a/distarray/dist/decorators.py +++ b/distarray/dist/decorators.py @@ -130,10 +130,10 @@ def process_return_value(self, context, result_key): client and return it. If all but one of the pulled values is None, return that non-None value only. """ - type_key = context._generate_key() - type_statement = "{} = str(type({}))".format(type_key, result_key) - context._execute(type_statement, targets=context.targets) - result_type_str = context._pull(type_key, targets=context.targets) + def get_type_str(key): + return str(type(key)) + result_type_str = context.apply(get_type_str, args=(result_key,), + targets=context.targets) def is_NoneType(typestring): return (typestring == "" or diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index 25294c25..fa079e07 100644 --- a/distarray/dist/distarray.py +++ b/distarray/dist/distarray.py @@ -245,9 +245,7 @@ def targets(self): def tondarray(self): """Returns the distributed array as an ndarray.""" arr = np.empty(self.shape, dtype=self.dtype) - local_name = self.context._generate_key() - self.context._execute('%s = %s.copy()' % (local_name, self.key), targets=self.targets) - local_arrays = self.context._pull(local_name, targets=self.targets) + local_arrays = self.get_localarrays() for local_array in local_arrays: maps = (list(ax_map.global_iter) for ax_map in local_array.distribution) @@ -323,11 +321,9 @@ def get_ndarrays(self): one ndarray per process """ - key = self.context._generate_key() - self.context._execute('%s = %s.get_localarray()' % (key, self.key), - targets=self.targets) - result = self.context._pull(key, targets=self.targets) - return result + def get(key): + return key.get_localarray() + return self.context.apply(get, args=(self.key,), targets=self.targets) def get_localarrays(self): """Pull the LocalArray objects from the engines. @@ -338,15 +334,14 @@ def get_localarrays(self): one localarray per process """ - result = self.context._pull(self.key, targets=self.targets) - return result + def get(key): + return key.copy() + return self.context.apply(get, args=(self.key,), targets=self.targets) def get_localshapes(self): - key = self.context._generate_key() - self.context._execute('%s = %s.local_shape' % (key, self.key), - targets=self.targets) - result = self.context._pull(key, targets=self.targets) - return result + def get(key): + return key.local_shape + return self.context.apply(get, args=(self.key,), targets=self.targets) # Binary operators