"""
This type stub file was generated by pyright.
"""

from collections import namedtuple
from collections.abc import MutableMapping

"""Columnar storage for DGLGraph."""
class Scheme(namedtuple('Scheme', ['shape', 'dtype'])):
    """The column scheme.

    Parameters
    ----------
    shape : tuple of int
        The feature shape.
    dtype : backend-specific type object
        The feature data type.
    """
    def __reduce__(self): # -> tuple[(shape: Unknown, dtype_str: Unknown) -> Self@Scheme, tuple[Unknown, Unknown]]:
        ...
    


def infer_scheme(tensor): # -> Scheme:
    """Infer column scheme from the given tensor data.

    Parameters
    ---------
    tensor : Tensor
        The tensor data.

    Returns
    -------
    Scheme
        The column scheme.
    """
    ...

class Column:
    """A column is a compact store of features of multiple nodes/edges.

    Currently, we use one dense tensor to batch all the feature tensors
    together (along the first dimension).

    Parameters
    ----------
    data : Tensor
        The initial data of the column.
    scheme : Scheme, optional
        The scheme of the column. Will be inferred if not provided.

    Attributes
    ----------
    data : Tensor
        The data of the column.
    scheme : Scheme
        The scheme of the column.
    """
    def __init__(self, data, scheme=...) -> None:
        ...
    
    def __len__(self):
        """The column length."""
        ...
    
    @property
    def shape(self):
        """Return the scheme shape (feature shape) of this column."""
        ...
    
    def __getitem__(self, idx):
        """Return the feature data given the index.

        Parameters
        ----------
        idx : utils.Index
            The index.

        Returns
        -------
        Tensor
            The feature data
        """
        ...
    
    def __setitem__(self, idx, feats): # -> None:
        """Update the feature data given the index.

        The update is performed out-placely so it can be used in autograd mode.
        For inplace write, please use ``update``.

        Parameters
        ----------
        idx : utils.Index or slice
            The index.
        feats : Tensor
            The new features.
        """
        ...
    
    def update(self, idx, feats, inplace): # -> None:
        """Update the feature data given the index.

        Parameters
        ----------
        idx : utils.Index
            The index.
        feats : Tensor
            The new features.
        inplace : bool
            If true, use inplace write.
        """
        ...
    
    def extend(self, feats, feat_scheme=...): # -> None:
        """Extend the feature data.

         Parameters
        ----------
        feats : Tensor
            The new features.
        feat_scheme : Scheme, optional
            The scheme
        """
        ...
    
    def clone(self): # -> Column:
        """Return a deepcopy of this column."""
        ...
    
    @staticmethod
    def create(data): # -> Column:
        """Create a new column using the given data."""
        ...
    
    def __repr__(self): # -> str:
        ...
    


class Frame(MutableMapping):
    """The columnar storage for node/edge features.

    The frame is a dictionary from feature fields to feature columns.
    All columns should have the same number of rows (i.e. the same first dimension).

    Parameters
    ----------
    data : dict-like, optional
        The frame data in dictionary. If the provided data is another frame,
        this frame will NOT share columns with the given frame. So any out-place
        update on one will not reflect to the other. The inplace update will
        be seen by both. This follows the semantic of python's container.
    num_rows : int, optional [default=0]
        The number of rows in this frame. If ``data`` is provided and is not empty,
        ``num_rows`` will be ignored and inferred from the given data.
    """
    def __init__(self, data=..., num_rows=...) -> None:
        ...
    
    def get_initializer(self, column=...): # -> (shape: Unknown, dtype: Unknown, ctx: Unknown, id_range: Unknown) -> Unknown | None:
        """Get the initializer for empty values for the given column.

        Parameters
        ----------
        column : str
            The column

        Returns
        -------
        callable
            The initializer
        """
        ...
    
    def set_initializer(self, initializer, column=...): # -> None:
        """Set the initializer for empty values, for a given column or all future
        columns.

        Initializer is a callable that returns a tensor given the shape and data type.

        Parameters
        ----------
        initializer : callable
            The initializer.
        column : str, optional
            The column name
        """
        ...
    
    def set_remote_init_builder(self, builder): # -> None:
        """Set an initializer builder to create a remote initializer for a new column to a frame.

        NOTE(minjie): This is a temporary solution. Will be replaced by KVStore in the future.

        The builder is a callable that returns an initializer. The returned initializer
        is also a callable that returns a tensor given a local tensor and tensor name.

        Parameters
        ----------
        builder : callable
            The builder to construct a remote initializer.
        """
        ...
    
    def get_remote_initializer(self, name): # -> None:
        """Get a remote initializer.

        NOTE(minjie): This is a temporary solution. Will be replaced by KVStore in the future.

        Parameters
        ----------
        name : string
            The column name.
        """
        ...
    
    @property
    def schemes(self): # -> dict[Unknown, Unknown | Scheme | None]:
        """Return a dictionary of column name to column schemes."""
        ...
    
    @property
    def num_columns(self): # -> int:
        """Return the number of columns in this frame."""
        ...
    
    @property
    def num_rows(self): # -> Unknown | int:
        """Return the number of rows in this frame."""
        ...
    
    def __contains__(self, name): # -> bool:
        """Return true if the given column name exists."""
        ...
    
    def __getitem__(self, name): # -> Column:
        """Return the column of the given name.

        Parameters
        ----------
        name : str
            The column name.

        Returns
        -------
        Column
            The column.
        """
        ...
    
    def __setitem__(self, name, data): # -> None:
        """Update the whole column.

        Parameters
        ----------
        name : str
            The column name.
        col : Column or data convertible to Column
            The column data.
        """
        ...
    
    def __delitem__(self, name): # -> None:
        """Delete the whole column.

        Parameters
        ----------
        name : str
            The column name.
        """
        ...
    
    def add_column(self, name, scheme, ctx): # -> None:
        """Add a new column to the frame.

        The frame will be initialized by the initializer.

        Parameters
        ----------
        name : str
            The column name.
        scheme : Scheme
            The column scheme.
        ctx : DGLContext
            The column context.
        """
        ...
    
    def add_rows(self, num_rows): # -> None:
        """Add blank rows to this frame.

        For existing fields, the rows will be extended according to their
        initializers.

        Parameters
        ----------
        num_rows : int
            The number of new rows
        """
        ...
    
    def update_column(self, name, data): # -> None:
        """Add or replace the column with the given name and data.

        Parameters
        ----------
        name : str
            The column name.
        data : Column or data convertible to Column
            The column data.
        """
        ...
    
    def append(self, other): # -> None:
        """Append another frame's data into this frame.

        If the current frame is empty, it will just use the columns of the
        given frame. Otherwise, the given data should contain all the
        column keys of this frame.

        Parameters
        ----------
        other : Frame or dict-like
            The frame data to be appended.
        """
        ...
    
    def clear(self): # -> None:
        """Clear this frame. Remove all the columns."""
        ...
    
    def __iter__(self): # -> Iterator[Unknown]:
        """Return an iterator of columns."""
        ...
    
    def __len__(self): # -> int:
        """Return the number of columns."""
        ...
    
    def keys(self): # -> _dict_keys[Unknown, Unknown] | _dict_keys[Unknown, Column]:
        """Return the keys."""
        ...
    
    def values(self): # -> _dict_values[Unknown, Unknown] | _dict_values[Column, Unknown]:
        """Return the values."""
        ...
    
    def clone(self): # -> Frame:
        """Return a clone of this frame.

        The clone frame does not share the underlying storage with this frame,
        i.e., adding or removing columns will not be visible to each other. However,
        they still share the tensor contents so any mutable operation on the column
        tensor are visible to each other. Hence, the function does not allocate extra
        tensor memory. Use :func:`~dgl.Frame.deepclone` for cloning
        a frame that does not share any data.

        Returns
        -------
        Frame
            A cloned frame.
        """
        ...
    
    def deepclone(self): # -> Frame:
        """Return a deep clone of this frame.

        The clone frame has an copy of this frame and any modification to the clone frame
        is not visible to this frame. The function allocate new tensors and copy the contents
        from this frame. Use :func:`~dgl.Frame.clone` for cloning a frame that does not
        allocate extra tensor memory.

        Returns
        -------
        Frame
            A deep-cloned frame.
        """
        ...
    


class FrameRef(MutableMapping):
    """Reference object to a frame on a subset of rows.

    Parameters
    ----------
    frame : Frame, optional
        The underlying frame. If not given, the reference will point to a
        new empty frame.
    index : utils.Index, optional
        The rows that are referenced in the underlying frame. If not given,
        the whole frame is referenced. The index should be distinct (no
        duplication is allowed).
    """
    def __init__(self, frame=..., index=...) -> None:
        ...
    
    @property
    def schemes(self): # -> dict[Unknown, Unknown | Scheme | None]:
        """Return the frame schemes.

        Returns
        -------
        dict of str to Scheme
            The frame schemes.
        """
        ...
    
    @property
    def num_columns(self): # -> int:
        """Return the number of columns in the referred frame."""
        ...
    
    @property
    def num_rows(self): # -> int:
        """Return the number of rows referred."""
        ...
    
    def set_initializer(self, initializer, column=...): # -> None:
        """Set the initializer for empty values.

        Initializer is a callable that returns a tensor given the shape and data type.

        Parameters
        ----------
        initializer : callable
            The initializer.
        column : str, optional
            The column name
        """
        ...
    
    def set_remote_init_builder(self, builder): # -> None:
        """Set an initializer builder to create a remote initializer for a new column to a frame.

        NOTE(minjie): This is a temporary solution. Will be replaced by KVStore in the future.

        The builder is a callable that returns an initializer. The returned initializer
        is also a callable that returns a tensor given a local tensor and tensor name.

        Parameters
        ----------
        builder : callable
            The builder to construct a remote initializer.
        """
        ...
    
    def get_initializer(self, column=...): # -> (shape: Unknown, dtype: Unknown, ctx: Unknown, id_range: Unknown) -> Unknown | None:
        """Get the initializer for empty values for the given column.

        Parameters
        ----------
        column : str
            The column

        Returns
        -------
        callable
            The initializer
        """
        ...
    
    def __contains__(self, name): # -> bool:
        """Return whether the column name exists."""
        ...
    
    def __iter__(self): # -> Iterator[Unknown]:
        """Return the iterator of the columns."""
        ...
    
    def __len__(self): # -> int:
        """Return the number of columns."""
        ...
    
    def keys(self): # -> _dict_keys[Unknown, Unknown] | _dict_keys[Unknown, Column]:
        """Return the keys."""
        ...
    
    def values(self): # -> _dict_values[Unknown, Unknown] | _dict_values[Column, Unknown]:
        """Return the values."""
        ...
    
    def __getitem__(self, key): # -> Self@FrameRef | LazyDict:
        """Get data from the frame.

        If the provided key is string, the corresponding column data will be returned.
        If the provided key is an index or a slice, the corresponding rows will be selected.
        The returned rows are saved in a lazy dictionary so only the real selection happens
        when the explicit column name is provided.

        Examples (using pytorch)
        ------------------------
        >>> # create a frame of two columns and five rows
        >>> f = Frame({'c1' : torch.zeros([5, 2]), 'c2' : torch.ones([5, 2])})
        >>> fr = FrameRef(f)
        >>> # select the row 1 and 2, the returned `rows` is a lazy dictionary.
        >>> rows = fr[Index([1, 2])]
        >>> rows['c1']  # only select rows for 'c1' column; 'c2' column is not sliced.

        Parameters
        ----------
        key : str or utils.Index
            The key.

        Returns
        -------
        Tensor or lazy dict or tensors
            Depends on whether it is a column selection or row selection.
        """
        ...
    
    def select_column(self, name):
        """Return the column of the given name.

        If only part of the rows are referenced, the fetching the whole column will
        also slice out the referenced rows.

        Parameters
        ----------
        name : str
            The column name.

        Returns
        -------
        Tensor
            The column data.
        """
        ...
    
    def select_rows(self, query): # -> LazyDict:
        """Return the rows given the query.

        Parameters
        ----------
        query : utils.Index or slice
            The rows to be selected.

        Returns
        -------
        utils.LazyDict
            The lazy dictionary from str to the selected data.
        """
        ...
    
    def __setitem__(self, key, val): # -> None:
        """Update the data in the frame. The update is done out-of-place.

        Parameters
        ----------
        key : str or utils.Index
            The key.
        val : Tensor or dict of tensors
            The value.

        See Also
        --------
        update
        """
        ...
    
    def update_data(self, key, val, inplace): # -> None:
        """Update the data in the frame.

        If the provided key is string, the corresponding column data will be updated.
        The provided value should be one tensor that have the same scheme and length
        as the column.

        If the provided key is an index, the corresponding rows will be updated. The
        value provided should be a dictionary of string to the data of each column.

        All updates are performed out-placely to be work with autograd. For inplace
        update, use ``update_column`` or ``update_rows``.

        Parameters
        ----------
        key : str or utils.Index
            The key.
        val : Tensor or dict of tensors
            The value.
        inplace: bool
            If True, update will be done in place
        """
        ...
    
    def update_column(self, name, data, inplace): # -> None:
        """Update the column.

        If this frameref spans the whole column of the underlying frame, this is
        equivalent to update the column of the frame.

        If this frameref only points to part of the rows, then update the column
        here will correspond to update part of the column in the frame. Raise error
        if the given column name does not exist.

        Parameters
        ----------
        name : str
            The column name.
        data : Tensor
            The update data.
        inplace : bool
            True if the update is performed inplacely.
        """
        ...
    
    def add_rows(self, num_rows): # -> None:
        """Add blank rows to the underlying frame.

        For existing fields, the rows will be extended according to their
        initializers.

        Note: only available for FrameRef that spans the whole column.  The row
        span will extend to new rows. Other FrameRefs referencing the same
        frame will not be affected.

        Parameters
        ----------
        num_rows : int
            Number of rows to add
        """
        ...
    
    def update_rows(self, query, data, inplace): # -> None:
        """Update the rows.

        If the provided data has new column, it will be added to the frame.

        See Also
        --------
        ``update_column``

        Parameters
        ----------
        query : utils.Index or slice
            The rows to be updated.
        data : dict-like
            The row data.
        inplace : bool
            True if the update is performed inplace.
        """
        ...
    
    def __delitem__(self, key): # -> None:
        """Delete data in the frame.

        If the provided key is a string, the corresponding column will be deleted.
        If the provided key is an index object or a slice, the corresponding rows will
        be deleted.

        Please note that "deleted" rows are not really deleted, but simply removed
        in the reference. As a result, if two FrameRefs point to the same Frame, deleting
        from one ref will not reflect on the other. However, deleting columns is real.

        Parameters
        ----------
        key : str or utils.Index
            The key.
        """
        ...
    
    def delete_rows(self, query): # -> None:
        """Delete rows.

        Please note that "deleted" rows are not really deleted, but simply removed
        in the reference. As a result, if two FrameRefs point to the same Frame, deleting
        from one ref will not reflect on the other. By contrast, deleting columns is real.

        Parameters
        ----------
        query : utils.Index
            The rows to be deleted.
        """
        ...
    
    def append(self, other): # -> None:
        """Append another frame into this one.

        Parameters
        ----------
        other : dict of str to tensor
            The data to be appended.
        """
        ...
    
    def clear(self): # -> None:
        """Clear the frame."""
        ...
    
    def is_contiguous(self): # -> bool:
        """Return whether this refers to a contiguous range of rows."""
        ...
    
    def is_span_whole_column(self): # -> bool:
        """Return whether this refers to all the rows."""
        ...
    
    def clone(self): # -> FrameRef:
        """Return a new reference to a clone of the underlying frame.

        Returns
        -------
        FrameRef
            A cloned frame reference.

        See Also
        --------
        dgl.Frame.clone
        """
        ...
    
    def deepclone(self): # -> FrameRef:
        """Return a new reference to a deep clone of the underlying frame.

        Returns
        -------
        FrameRef
            A deep-cloned frame reference.

        See Also
        --------
        dgl.Frame.deepclone
        """
        ...
    


def frame_like(other, num_rows=...): # -> Frame:
    """Create an empty frame that has the same initializer as the given one.

    Parameters
    ----------
    other : Frame
        The given frame.
    num_rows : int
        The number of rows of the new one. If None, use other.num_rows
        (Default: None)

    Returns
    -------
    Frame
        The new frame.
    """
    ...

def sync_frame_initializer(new_frame, reference_frame): # -> None:
    """Set the initializers of the new_frame to be the same as the reference_frame,
    for both the default initializer and per-column initializers.

    Parameters
    ----------
    new_frame : Frame
        The frame to set initializers
    reference_frame : Frame
        The frame to copy initializers
    """
    ...

