
    h6                        S SK Jr  S SKJrJr  S SKJr  \(       a  S SKJ	r	J
r
Jr  \S   r " S S5      r " S S	5      rg)
    )annotations)TYPE_CHECKINGUnionN)	DataFrame	LazyFrameSeries)r   r   c                      \ rS rSr% SrS\S'   S\S'   S\S'   S\S'   S\S	'   S
\S'   SS jr\SS j5       r\SS j5       r	SS jr
SS jrSS jrSS jrSrg)PolarsSlice   zq
Apply Python slice object to Polars DataFrame or Series.

Has full support for negative indexing and/or stride.
intstopstartstrideslice_lengthboolis_unboundedFrameOrSeriesobjc                    Xl         g Nr   selfr   s     M/home/james-whalen/.local/lib/python3.13/site-packages/polars/_utils/slice.py__init__PolarsSlice.__init__           c                    U R                  5       n[        U[        R                  5      (       a  U$ UR	                  5       $ )z.Return lazy variant back to its original type.)collect
isinstanceplr   	to_series)lazyoriginalframes      r   _as_originalPolarsSlice._as_original   s1     "8R\\::uQ@QQr   c                    [        U [        R                  5      (       a  U R                  5       R	                  5       $ U R	                  5       $ )z2Make lazy to ensure efficient/consistent handling.)r!   r"   r   to_framer$   r   s    r   _lazifyPolarsSlice._lazify$   s4     )33		(B(Bs||~""$R
Rr   c                    UR                  U R                  U R                  5      R                  U R                  5      $ )z&Logic for slices with positive stride.)slicer   r   gather_everyr   r   s     r   _slice_positivePolarsSlice._slice_positive)   s.     yyT%6%67DDT[[QQr   c                    [        U R                  5      nUR                  U R                  S-   U R                  5      R                  5       nUS:  a  UR                  U5      $ U$ )z&Logic for slices with negative stride.   )absr   r.   r   r   reverser/   )r   r   r   	lazyslices       r   _slice_negativePolarsSlice._slice_negative.   sS    T[[!IIdii!mT->->?GGI	281*y%%f-L9Lr   c                   [        U R                  5      n[        UR                  UR                  UR
                  5      R                  U5      u  p4nUS:  a  US:*  =(       a    XB:  U l        OUS:H  =(       a    X2S-
  :  U l        U R                  R                  5       (       a  SU l	        OGU R                  (       a  X l	        O/X4:X  d  US:  a  X4:  d  US:  a  X4:  a  SO[        XC-
  5      U l	        X4UsU l        U l        U l        g)zENormalise slice bounds, identify unbounded and/or zero-length slices.r3   r   N)lenr   r.   r   r   stepindicesr   is_emptyr   r4   r   )r   sobj_lenr   r   r   s         r   _slice_setupPolarsSlice._slice_setup4   s     dhh-#AGGQVVQVV<DDWMV Q;!&! B$/D!% G%Q;2FD 88 !D '
 ]
u|
u|	  &  .3&*
DIt{r   c                d   U R                  U5        U R                  S:X  a  U R                  R                  5       $ U R                  (       aT  U R
                  S;   aD  U R
                  S:  a  U R                  R                  5       $ U R                  R                  5       $ U R                  S:  aP  U R                  S:  a@  U R
                  S:X  a0  U R                  R                  U R                  U R                  5      $ U R
                  S:  a9  U R                  S:X  a)  U R                  R                  U R                  S-   S5      $ U R                  U R                  5      nU R
                  S:  a  U R                  U5      OU R                  U5      nU R                  X0R                  5      $ )zFApply a slice operation, taking advantage of any potential fast paths.r   )r:   r3   r3   )rA   r   r   clearr   r   r5   cloner   r   r.   r+   r0   r7   r'   )r   r?   lazyobjsliceds       r   applyPolarsSlice.applyQ   sC    	! !88>>##4;;'#9*.++/488##%P@PPZZ1_aDKK14D88>>$**d.?.?@@[[1_!2!2a!788>>$))a-33 ll488,G ;;? $$W-))'2 
 $$VXX66r   )r   r   r   r   r   r   N)r   r   returnNone)r$   r   r%   r   rJ   r   )r   r   rJ   r   )r   r   rJ   r   )r?   r.   rJ   rK   )r?   r.   rJ   r   )__name__
__module____qualname____firstlineno____doc____annotations__r   staticmethodr'   r+   r0   r7   rA   rH   __static_attributes__ r   r   r
   r
      sr     IJK	 R R
 S SR
MA:7r   r
   c                  8    \ rS rSr% SrS\S'   S	S jrS
S jrSrg)LazyPolarsSlicem   z
Apply python slice object to Polars LazyFrame.

Only slices with efficient computation paths that map directly
to existing lazy methods are supported.
r   r   c                    Xl         g r   r   r   s     r   r   LazyPolarsSlice.__init__w   r   r   c                0   UR                   =(       d    SnUR                  =(       d    SnUR                  (       a  UR                  S:  a  Sn[        U5      eUS:  aL  US:  d  UR                  b9  X!R                  :w  a*  USs=:  a  U:  a  O  OUR                  b  Sn[        U5      eUS:  a  UR                  b  X!R                  :  dD  US:  aX  UR                   bK  UR                  b>  UR                  UR                   s=:  a  S:  a  O  OU R                  R                  5       $ UR                   c  UR                  c  US:X  a  U R                  R                  5       $ US:  a  U R                  R                  U5      $ US:X  a  U R                  R                  5       $ US:  a2  U R                  R                  5       R                  [        U5      5      $ GOUSs=:  a  U:  aq  O  OnUR                  ca  U R                  R                  UR                   S-   5      R                  5       n[        U5      S:X  a  U$ UR                  [        U5      5      $ US:X  aW  UR                  =(       d    SS:  a>  U R                  R                  UR                  5      nUS:X  a  U$ UR                  U5      $ US:  aP  UR                  cC  US:  a=  U R                  R                  [        U5      5      nUS:X  a  U$ UR                  U5      $ US:  ao  UR                  b  UR                  S:  aR  UR                  c  SOUR                  U-
  nU R                  R                  X&5      nUS:X  a  U$ UR                  U5      $ SU< S3n[        U5      e)	z
Apply a slice operation.

Note that LazyFrame is designed primarily for efficient computation and does not
know its own length so, unlike DataFrame, certain slice patterns (such as those
requiring negative stop/step) may not be supported.
r   r3   z.negative stop is not supported for lazy slicesNz?negative stride is not supported in conjunction with start+stopr:   zthe given slice zv is not supported by lazy computation

Consider a more efficient approach, or construct explicitly with other methods.)r   r<   r   
ValueErrorr   rD   rE   r/   r5   r4   headtailr.   )r   r?   r   r<   msgr   r   s          r   rH   LazyPolarsSlice.applyz   s    1vv{ 66affqjBCS/!!8aff&8uA$$W o% 1H!&&,&&1H$);!''@VUV@V88>>## WW_qyxx~~''xx,,T22xx''))xx'')66s4yAA  a$166>((--!,446Ct9>3K0@0@T0KK aZQVV[qQ.((--'C193A3+;+;D+AA QY166>dQh((--E
+C193A3+;+;D+AA QYAFFNaffk$%FFN4!&&5.L((..5C193A3+;+;D+AA qe $b b 	 or   r   N)r   r   rJ   rK   )r?   r.   rJ   r   )	rL   rM   rN   rO   rP   rQ   r   rH   rS   rT   r   r   rV   rV   m   s     
N]r   rV   )
__future__r   typingr   r   polars._reexport	_reexportr"   polarsr   r   r   r   r
   rV   rT   r   r   <module>re      s;    " ' 33/0M]7 ]7@j jr   