Skip to content

Commit 24d45e2

Browse files
eteqeteq
committed
fixes to support array operations of Quantity objects
1 parent 35ad30e commit 24d45e2

File tree

1 file changed

+58
-56
lines changed

1 file changed

+58
-56
lines changed

‎astropy/units/quantity.py‎

Lines changed: 58 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -190,91 +190,93 @@ def cgs(self):
190190

191191
# Arithmetic operations
192192
def __add__(self, other):
193-
""" Addition between `Quantity` objects. All operations return a new `Quantity` object
194-
with the units of the **left** object.
193+
""" Addition between `Quantity` objects and other objects. If
194+
they are both `Quantity` objects, results in the units of the
195+
**left** object if they are compatible, otherwise this fails.
195196
"""
196-
if not isinstance(other, Quantity):
197-
raise TypeError("Object of type '{0}' cannot be added with a Quantity object. Addition is only supported between Quantity objects.".format(other.__class__))
198-
return Quantity(self.value + other.to(self.unit).value, unit=self.unit)
197+
if isinstance(other, Quantity):
198+
return Quantity(self.value + other.to(self.unit).value, unit=self.unit)
199+
else:
200+
raise TypeError("Object of type '{0}' cannot be added with a "
201+
"Quantity object. Addition is only supported between Quantity "
202+
"objects with compatible units.".format(other.__class__))
203+
199204

200205
def __sub__(self, other):
201-
""" Subtraction between `Quantity` objects. All operations return a new `Quantity` object
202-
with the units of the **left** object.
206+
""" Subtraction between `Quantity` objects and other objects.
207+
If they are both `Quantity` objects, results in the units of the
208+
**left** object if they are compatible, otherwise this fails.
203209
"""
204-
if not isinstance(other, Quantity):
205-
raise TypeError("Object of type '{0}' cannot be subtracted with a Quantity object. Subtraction is only supported between Quantity objects.".format(other.__class__))
206-
return Quantity(self.value - other.to(self.unit).value, unit=self.unit)
210+
if isinstance(other, Quantity):
211+
return Quantity(self.value - other.to(self.unit).value, unit=self.unit)
212+
else:
213+
raise TypeError("Object of type '{0}' cannot be added with a "
214+
"Quantity object. Addition is only supported between Quantity "
215+
"objects with compatible units.".format(other.__class__))
216+
207217

208218
def __mul__(self, other):
209-
""" Multiplication between `Quantity` objects or numbers with `Quantity` objects. For operations between two `Quantity` instances,
210-
returns a new `Quantity` object with the units of the **left** object.
219+
""" Multiplication between `Quantity` objects and other objects.
211220
"""
212221
if isinstance(other, Quantity):
213-
return Quantity(self.value * other.value, unit=self.unit*other.unit)
214-
215-
elif isinstance(other, numbers.Number):
216-
return Quantity(other*self.value, unit=self.unit)
217-
222+
return Quantity(self.value * other.value, unit=self.unit * other.unit)
218223
elif isinstance(other, UnitBase):
219-
return Quantity(self.value, unit=self.unit*other)
220-
224+
return Quantity(self.value, unit=other * self.unit)
221225
else:
222-
raise TypeError("Object of type '{0}' cannot be multiplied with a Quantity object.".format(other.__class__))
223-
226+
try:
227+
return Quantity(other * self.value, unit=self.unit)
228+
except TypeError:
229+
raise TypeError("Object of type '{0}' cannot be multiplied with a Quantity object.".format(other.__class__))
224230

225231
def __rmul__(self, other):
226-
""" Right multiplication between `Quantity` object and a number. """
227-
228-
if isinstance(other, numbers.Number):
229-
return Quantity(other*self.value, unit=self.unit)
230-
231-
elif isinstance(other, UnitBase):
232-
return Quantity(self.value, unit=self.unit*other)
233-
234-
else:
235-
raise TypeError("Object of type '{0}' cannot be multiplied with a Quantity object.".format(other.__class__))
232+
""" Right Multiplication between `Quantity` objects and other
233+
objects.
234+
"""
235+
return self.__mul__(other)
236236

237237
def __div__(self, other):
238-
""" Division between `Quantity` objects. This operation returns a dimensionless object. """
238+
""" Division between `Quantity` objects and other objects.
239+
"""
239240
if isinstance(other, Quantity):
240-
return Quantity(self.value / other.value, unit=self.unit/other.unit)
241-
242-
elif isinstance(other, numbers.Number):
243-
return Quantity(self.value / other, unit=self.unit)
244-
241+
return Quantity(self.value / other.value, unit=self.unit / other.unit)
245242
elif isinstance(other, UnitBase):
246-
return Quantity(self.value, unit=self.unit/other)
247-
243+
return Quantity(self.value, unit=self.unit / other)
244+
elif isinstance(other, numbers.Number) and hasattr(self.unit, "bases"):
245+
new_unit_bases = copy.copy(self.unit.bases)
246+
new_unit_powers = [-p for p in self.unit.powers]
247+
return Quantity(other / self.value, unit=CompositeUnit(1., new_unit_bases, new_unit_powers))
248248
else:
249-
raise TypeError("Object of type '{0}' cannot be divided with a Quantity object.".format(other.__class__))
249+
try:
250+
return Quantity(self.value / other, unit=self.unit)
251+
except TypeError:
252+
raise TypeError("Object of type '{0}' cannot be diveded with a Quantity object.".format(other.__class__))
250253

251254
def __rdiv__(self, other):
252-
""" Division between `Quantity` objects. This operation returns a dimensionless object. """
253-
if isinstance(other, numbers.Number):
254-
if hasattr(self.unit, "bases"):
255-
new_unit_bases = copy.copy(self.unit.bases)
256-
new_unit_powers = [-p for p in self.unit.powers]
257-
return Quantity(other / self.value, unit=CompositeUnit(1., new_unit_bases, new_unit_powers))
258-
else:
259-
return Quantity(other / self.value, unit=1./self.unit)
260-
255+
""" Right Division between `Quantity` objects and other objects.
256+
"""
257+
if isinstance(other, Quantity):
258+
return Quantity(other.value / self.value, unit=other.unit / self.unit)
261259
elif isinstance(other, UnitBase):
262-
return Quantity(1./self.value, unit=other/self.unit)
263-
260+
return Quantity(1. / self.value, unit=other / self.unit)
264261
else:
265-
raise TypeError("Object of type '{0}' cannot be divided with a Quantity object.".format(other.__class__))
262+
try:
263+
return Quantity(other / self.value, unit=1. / self.unit)
264+
except TypeError:
265+
raise TypeError("Object of type '{0}' cannot be diveded with a Quantity object.".format(other.__class__))
266266

267267
def __truediv__(self, other):
268-
""" Division between `Quantity` objects. This operation returns a dimensionless object. """
268+
""" Division between `Quantity` objects. """
269269
return self.__div__(other)
270270

271271
def __rtruediv__(self, other):
272-
""" Division between `Quantity` objects. This operation returns a dimensionless object. """
272+
""" Division between `Quantity` objects. """
273273
return self.__rdiv__(other)
274274

275275
def __pow__(self, p):
276-
""" Raise quantity object to a power. """
277-
return Quantity(self.value**p, unit=self.unit**p)
276+
""" Raise `Quantity` object to a power. """
277+
if hasattr(p, 'unit'):
278+
raise TypeError('Cannot raise a Quantity object to a power of something with a unit')
279+
return Quantity(self.value ** p, unit=self.unit ** p)
278280

279281
# Comparison operations
280282
def __eq__(self, other):

0 commit comments

Comments
 (0)