Exemple #1
0
def rewrite_slice(node, nopython):
    """
    Rewrites array slices to its native equivalent without
    using the Python API.

        node:       ast.Subscript with an array type as result
        nopython:   whether the node is encountered in a nopython context
    """
    # assert self.nopython

    if isinstance(node.slice, ast.ExtSlice):
        dims = node.slice.dims
    else:
        assert not isinstance(node.slice, ast.Ellipsis)
        dims = [node.slice]

    slices = []
    src_dim = 0
    dst_dim = 0

    all_slices = True
    for subslice in dims:
        slices.append(create_slice_dim_node(subslice, src_dim, dst_dim))

        if subslice.type.is_slice:
            src_dim += 1
            dst_dim += 1
        elif nodes.is_newaxis(subslice):
            all_slices = False
            dst_dim += 1
        else:
            assert subslice.type.is_int
            all_slices = False
            src_dim += 1

    #if all_slices and all(empty(subslice) for subslice in slices):
    #    return node.value

    # print node, node.type
    return NativeSliceNode(node.type, node.value, slices, nopython)
Exemple #2
0
def unellipsify(node, slices, subscript_node):
    """
    Given an array node `node`, process all AST slices and create the
    final type:

        - process newaxes (None or numpy.newaxis)
        - replace Ellipsis with a bunch of ast.Slice objects
        - process integer indices
        - append any missing slices in trailing dimensions
    """
    type = node.variable.type

    if not type.is_array:
        assert type.is_object
        return object_, node

    if (len(slices) == 1 and nodes.is_constant_index(slices[0]) and
            slices[0].value.pyval is Ellipsis):
        # A[...]
        return type, node

    result = []
    seen_ellipsis = False

    # Filter out newaxes
    newaxes = [newaxis for newaxis in slices if nodes.is_newaxis(newaxis)]
    n_indices = len(slices) - len(newaxes)

    full_slice = ast.Slice(lower=None, upper=None, step=None)
    full_slice.variable = Variable(typesystem.SliceType())
    ast.copy_location(full_slice, slices[0])

    # process ellipses and count integer indices
    indices_seen = 0
    for slice_node in slices[::-1]:
        slice_type = slice_node.variable.type
        if slice_type.is_ellipsis:
            if seen_ellipsis:
                result.append(full_slice)
            else:
                nslices = type.ndim - n_indices + 1
                result.extend([full_slice] * nslices)
                seen_ellipsis = True
        elif (slice_type.is_slice or slice_type.is_int or
              nodes.is_newaxis(slice_node)):
            indices_seen += slice_type.is_int
            result.append(slice_node)
        else:
            # TODO: Coerce all object operands to integer indices?
            # TODO: (This will break indexing with the Ellipsis object or
            # TODO:  with slice objects that we couldn't infer)
            return object_, nodes.CoercionNode(node, object_)

    # Reverse our reversed processed list of slices
    result.reverse()

    # append any missing slices (e.g. a2d[:]
    result_length = len(result) - len(newaxes)
    if result_length < type.ndim:
        nslices = type.ndim - result_length
        result.extend([full_slice] * nslices)

    subscript_node.slice = ast.ExtSlice(result)
    ast.copy_location(subscript_node.slice, slices[0])

    # create the final array type and set it in value.variable
    result_dtype = node.variable.type.dtype
    result_ndim = node.variable.type.ndim + len(newaxes) - indices_seen
    if result_ndim > 0:
        result_type = result_dtype[(slice(None),) * result_ndim]
    elif result_ndim == 0:
        result_type = result_dtype
    else:
        result_type = object_

    return result_type, node
Exemple #3
0
def unellipsify(node, slices, subscript_node):
    """
    Given an array node `node`, process all AST slices and create the
    final type:

        - process newaxes (None or numpy.newaxis)
        - replace Ellipsis with a bunch of ast.Slice objects
        - process integer indices
        - append any missing slices in trailing dimensions
    """
    type = node.variable.type

    if not type.is_array:
        assert type.is_object
        return object_, node

    if (len(slices) == 1 and nodes.is_constant_index(slices[0]) and
            slices[0].value.pyval is Ellipsis):
        # A[...]
        return type, node

    result = []
    seen_ellipsis = False

    # Filter out newaxes
    newaxes = [newaxis for newaxis in slices if nodes.is_newaxis(newaxis)]
    n_indices = len(slices) - len(newaxes)

    full_slice = ast.Slice(lower=None, upper=None, step=None)
    full_slice.variable = Variable(typesystem.slice_)
    ast.copy_location(full_slice, slices[0])

    # process ellipses and count integer indices
    indices_seen = 0
    for slice_node in slices[::-1]:
        slice_type = slice_node.variable.type
        if slice_type.is_ellipsis:
            if seen_ellipsis:
                result.append(full_slice)
            else:
                nslices = type.ndim - n_indices + 1
                result.extend([full_slice] * nslices)
                seen_ellipsis = True
        elif (slice_type.is_slice or slice_type.is_int or
              nodes.is_newaxis(slice_node)):
            indices_seen += slice_type.is_int
            result.append(slice_node)
        else:
            # TODO: Coerce all object operands to integer indices?
            # TODO: (This will break indexing with the Ellipsis object or
            # TODO:  with slice objects that we couldn't infer)
            return object_, nodes.CoercionNode(node, object_)

    # Reverse our reversed processed list of slices
    result.reverse()

    # append any missing slices (e.g. a2d[:]
    result_length = len(result) - len(newaxes)
    if result_length < type.ndim:
        nslices = type.ndim - result_length
        result.extend([full_slice] * nslices)

    subscript_node.slice = ast.ExtSlice(result)
    ast.copy_location(subscript_node.slice, slices[0])

    # create the final array type and set it in value.variable
    result_dtype = node.variable.type.dtype
    result_ndim = node.variable.type.ndim + len(newaxes) - indices_seen
    if result_ndim > 0:
        result_type = result_dtype[(slice(None),) * result_ndim]
    elif result_ndim == 0:
        result_type = result_dtype
    else:
        result_type = object_

    return result_type, node