test_ytarray.py 文件源码

def test_subclass():

    class YTASubclass(YTArray):
        pass

    a = YTASubclass([4, 5, 6], 'g')
    b = YTASubclass([7, 8, 9], 'kg')
    nu = YTASubclass([10, 11, 12], '')
    nda = np.array([3, 4, 5])
    yta = YTArray([6, 7, 8], 'mg')
    loq = [YTQuantity(6, 'mg'), YTQuantity(7, 'mg'), YTQuantity(8, 'mg')]
    ytq = YTQuantity(4, 'cm')
    ndf = np.float64(3)

    def op_comparison(op, inst1, inst2, compare_class):
        assert_isinstance(op(inst1, inst2), compare_class)
        assert_isinstance(op(inst2, inst1), compare_class)

    ops = [operator.mul, operator.truediv]
    if hasattr(operator, "div"):
        ops.append(operator.div)
    for op in ops:
        for inst in (b, ytq, ndf, yta, nda, loq):
            op_comparison(op, a, inst, YTASubclass)

        op_comparison(op, ytq, nda, YTArray)
        op_comparison(op, ytq, yta, YTArray)

    for op in (operator.add, operator.sub):
        op_comparison(op, nu, nda, YTASubclass)
        op_comparison(op, a, b, YTASubclass)
        op_comparison(op, a, yta, YTASubclass)
        op_comparison(op, a, loq, YTASubclass)

    assert_isinstance(a[0], YTQuantity)
    assert_isinstance(a[:], YTASubclass)
    assert_isinstance(a[:2], YTASubclass)
    assert_isinstance(YTASubclass(yta), YTASubclass)