
    bZhX                        d Z ddlmZ ddlmZ ddlmZ ddlmZm	Z	 ddl
mZmZmZ er,ddlmZmZ ddlmZmZ dd	lZdd	lZdd
lmZ ddlmZ eeef   ZneZdgZ G d de      Zy	)zxSchema.

Adapted from Polars implementation at:
https://github.com/pola-rs/polars/blob/main/py-polars/polars/schema.py.
    )annotations)OrderedDict)partial)TYPE_CHECKINGcast)ImplementationVersionparse_version)IterableMapping)AnyClassVarN)DType)DTypeBackendSchemac                       e Zd ZU dZej
                  Zded<   	 d	 	 	 d fdZddZ	ddZ
ddZddZ	 d	 	 	 dd	Zdd
Z xZS )r   a<  Ordered mapping of column names to their data type.

    Arguments:
        schema: The schema definition given by column names and their associated
            *instantiated* Narwhals data type. Accepts a mapping or an iterable of tuples.

    Examples:
        Define a schema by passing *instantiated* data types.

        >>> import narwhals as nw
        >>> schema = nw.Schema({"foo": nw.Int8(), "bar": nw.String()})
        >>> schema
        Schema({'foo': Int8, 'bar': String})

        Access the data type associated with a specific column name.

        >>> schema["foo"]
        Int8

        Access various schema properties using the `names`, `dtypes`, and `len` methods.

        >>> schema.names()
        ['foo', 'bar']
        >>> schema.dtypes()
        [Int8, String]
        >>> schema.len()
        2
    zClassVar[Version]_versionc                0    |xs i }t         |   |       y N)super__init__)selfschema	__class__s     ^/mounts/lovelace/software/anaconda3/envs/py312/lib/python3.12/site-packages/narwhals/schema.pyr   zSchema.__init__A   s     2     c                4    t        | j                               S )zXGet the column names of the schema.

        Returns:
            Column names.
        )listkeysr   s    r   nameszSchema.namesG   s     DIIK  r   c                4    t        | j                               S )z^Get the data types of the schema.

        Returns:
            Data types of schema.
        )r   valuesr    s    r   dtypeszSchema.dtypesO   s     DKKM""r   c                    t        |       S )zbGet the number of columns in the schema.

        Returns:
            Number of columns.
        )lenr    s    r   r&   z
Schema.lenW   s     4yr   c                n     ddl }ddlm  |j                   fd j	                         D              S )a7  Convert Schema to a pyarrow Schema.

        Returns:
            A pyarrow Schema.

        Examples:
            >>> import narwhals as nw
            >>> schema = nw.Schema({"a": nw.Int64(), "b": nw.Datetime("ns")})
            >>> schema.to_arrow()
            a: int64
            b: timestamp[ns]
        r   Nnarwhals_to_native_dtypec              3  N   K   | ]  \  }}| |j                         f  y wr   r   ).0namedtyper)   r   s      r   	<genexpr>z"Schema.to_arrow.<locals>.<genexpr>p   s.      
e +E4==AB
s   "%)pyarrownarwhals._arrow.utilsr)   r   items)r   par)   s   ` @r   to_arrowzSchema.to_arrow_   s1     	Bryy 
#zz|
 
 	
r   c                   ddl }ddlm} t        |t        j
                  t        |      | j                        }|t        |t              r,| j                         D ci c]  \  }}| |||       c}}S t        |      }t        |      t        |       k7  rnddlm}m}	m}
 t        |      t        |       }}t         |	|j#                   |	 |
|      |            |            }d|d|d	| d
|d    d| d}t%        |      t'        | j)                         | j+                         |      D ci c]  \  }}}| |||       c}}}S c c}}w c c}}}w )a  Convert Schema to an ordered mapping of column names to their pandas data type.

        Arguments:
            dtype_backend: Backend(s) used for the native types. When providing more than
                one, the length of the iterable must be equal to the length of the schema.

        Returns:
            An ordered mapping of column names to their pandas data type.

        Examples:
            >>> import narwhals as nw
            >>> schema = nw.Schema({"a": nw.Int64(), "b": nw.Datetime("ns")})
            >>> schema.to_pandas()
            {'a': 'int64', 'b': 'datetime64[ns]'}

            >>> schema.to_pandas("pyarrow")
            {'a': 'Int64[pyarrow]', 'b': 'timestamp[ns][pyarrow]'}
        r   Nr(   )implementationbackend_versionversion)r.   dtype_backend)chainislicerepeatz	Provided z) `dtype_backend`(s), but schema contains z1 field(s).
Hint: instead of
    schema.to_pandas(z+)
you may want to use
    schema.to_pandas(z)
or
    schema.to_pandas())pandasnarwhals._pandas_like.utilsr)   r   r   PANDASr
   r   
isinstancestrr2   tupler&   	itertoolsr:   r;   r<   from_iterable
ValueErrorzipr   r#   )r   r9   pdr)   to_native_dtyper-   r.   backendsr:   r;   r<   n_usern_actual
suggestionmsgbackends                   r   	to_pandaszSchema.to_pandasu   su   * 	H!$)00)"-MM	
  J}c$B $(::<D% oEOO 
 ]+H8}D	);;#&x=#d)"++F6(3CX,NOQY
  z)RS[R^ _,,4: 6,,4QK= 9,,6<q:  !o% -0		T[[]H,U (D% oEII 32s    E6Ec                     ddl }ddlm t        |       fd j	                         D        }dk\  r |j
                  |      S t        dt        |            S )ax  Convert Schema to a polars Schema.

        Returns:
            A polars Schema or plain dict (prior to polars 1.0).

        Examples:
            >>> import narwhals as nw
            >>> schema = nw.Schema({"a": nw.Int64(), "b": nw.Datetime("ns")})
            >>> schema.to_polars()
            Schema({'a': Int64, 'b': Datetime(time_unit='ns', time_zone=None)})
        r   Nr(   c              3  R   K   | ]  \  }}| |j                          f   yw))r7   Nr+   )r,   r-   r.   r)   
pl_versionr   s      r   r/   z#Schema.to_polars.<locals>.<genexpr>   s7      
 e (4==*
s   $')   r   r   	pl.Schema)polarsnarwhals._polars.utilsr)   r
   r2   r   r   dict)r   plr   r)   rS   s   `  @@r   	to_polarszSchema.to_polars   s`     	C"2&

  $zz|
 Y& BIIf	
 k4<0	
r   r   )r   z8Mapping[str, DType] | Iterable[tuple[str, DType]] | NonereturnNone)r[   z	list[str])r[   zlist[DType])r[   int)r[   z	pa.Schema)r9   z%DTypeBackend | Iterable[DTypeBackend]r[   zdict[str, Any])r[   rU   )__name__
__module____qualname____doc__r	   MAINr   __annotations__r   r!   r$   r&   r4   rP   rZ   __classcell__)r   s   @r   r   r   !   sk    : #*,,H. RV!N!	!!#
. FJ<B<	<|
r   )ra   
__future__r   collectionsr   	functoolsr   typingr   r   narwhals._utilsr   r	   r
   collections.abcr   r   r   r   rV   rY   r0   r3   narwhals.dtypesr   narwhals.typingr   rB   
BaseSchema__all__r    r   r   <module>rp      s]    # #  & B B1$%,S%Z(J J*p
Z p
r   