
    <i'                     J   d dl Z d dlmZ d dlmZmZmZmZ d dlm	Z
 d dlmZ d dlmZmZmZmZ d dlmZmZ d dlmZmZmZ d d	lmZ d
ddddZdedefdZ edd          Zdee         de de dee         fdZ!dededededede"de de dee         fd Z# G d! d"          Z$dS )#    N)
namedtuple)ListAnyOptionalUnion)_msgs)extract_args)commandKeyFloatCommandItem)SimpleErrorDatabase)distance
geo_encode
geo_decode)ZSetgMbP?gp\D?gD
)?
@   )kmmiftmunit_argreturnc                     t                               |                                                                           }|t	          t
          j                  |S )zvnumber of meters in a unit.
    :param unit_arg: unit name (km, mi, ft, m)
    :returns: number of meters in unit
    )	UNIT_TO_Mgetdecodelowerr   msgsGEO_UNSUPPORTED_UNIT)r   units     /Users/kimhansen/Desktop/03 Workspace/ceo-agents/chl-effectiveness/mcp-servers/whoop/.venv/lib/python3.11/site-packages/fakeredis/commands_mixins/geo_mixin.pytranslate_meters_to_unitr$      sF    
 ==**002233D|$3444K    	GeoResultzname long lat hash distanceitems	withcoordwithdistc                 t   g }| D ]}|j         g}|r-|                    t          j        |j        d                     |rG|                    t          j        |j        d          t          j        |j        d          g           t          |          dk    r|d         }|                    |           |S )zParse list of GeoResults to redis response
    :param withcoord: include coordinates in response
    :param withdist: include distance in response
    :returns: Parsed list
    Fr   r   )nameappendr   encoder   longlatlen)r'   r(   r)   resitemnew_items         r#   _parse_resultsr4      s     C  I; 	@OOEL>>??? 	]OOU\$)U;;U\$(TY=Z=Z[\\\x==A{H

8Jr%   zsetr/   r.   radiusconvcount	count_anydescc           
      X   g }|                                  D ]s\  }	}
t          |
          \  }}}}t          ||f||f          |z  }||k     r=|                    t	          |	|||
|                     |rt          |          |k    r ntt          |d |          }|r
|d|         }|S )a  Find items within area (lat,long)+radius
    :param zset: list of items to check
    :param lat: latitude
    :param long: longitude
    :param radius: radius in whatever units
    :param conv: conversion of radius to meters
    :param count: number of results to give
    :param count_any: should we return any results that match? (vs. sorted)
    :param desc: should results be sorted descending order?
    :returns: List of GeoResults
    c                     | j         S N)r   )xs    r#   <lambda>z_find_near.<locals>.<lambda>M   s    AJ r%   )keyreverseN)r'   r   r   r,   r&   r0   sorted)r5   r/   r.   r6   r7   r8   r9   r:   resultsr+   _hashp_latp_long_dists                  r#   
_find_nearrI   0   s    * Gzz||  e(//vq!#t55<&==NN9T65%FFGGG S\\U22W"6"6EEEG "&5&/Nr%   c                       e Zd Zdededdf fdZdedee         dede	fd	Z
d
edededededeee         e	f         fdZ ed ee          fef          d
edede	fd            Z ed ee          efef          d
ededeeedf                  fd            Z ed ee          efef          d
ededeeee                           fd            Z ed ee          eefef          d
ededededee         f
d            Z ed ee          eeefef          d
edededededee         fd            Z ed ee          eeefef          d
edededededeee         e	f         fd            Z ed ee          eefef          d
ed edededee         f
d!            Z ed" ee          eefef          d
ed edededee         f
d#            Z ed$ ee          fef          d
ededee         fd%            Z ed&e ee          fef          d'ed(ededee         fd)            Z xZS )*GeoCommandsMixinargskwargsr   Nc                 L     t          t          |           j        |i | |  d S r=   )superrK   __init__)selfrL   rM   	__class__s      r#   rP   zGeoCommandsMixin.__init__T   s2    .%%.????r%   	item_namegeo_results	scoredistc                 P   t          || j        | j                            |          t                                }t                      |_        |D ]2}|r|j        n|j        }|j                            |j        |           3|	                                 t          |          S )N)r2   default)r   _dbr   r   valuer   hashaddr+   	writebackr0   )rQ   rS   rT   rU   db_itemr2   vals          r#   _store_geo_resultsz#GeoCommandsMixin._store_geo_resultsX   s    iY8O8OY]Y_Y_``` 	. 	.D#,;$--$)CMdi----;r%   r@   r.   r/   r6   c           
         t          |ddd          \  \  }}}}	}
}}}}|	pt          j        }	t          |          dk    rt	          |d                   nd}t          |j        |||||	|
|          }|r'|                     ||d           t          |          S |r'|                     ||d           t          |          S t          |||          S )N)r(   r)   withhash+countanyr:   z*storez
*storedistFerror_on_unexpectedleft_from_first_unexpectedr   r   )rU   T)	r	   sysmaxsizer0   r$   rI   rY   r_   r4   )rQ   r@   r.   r/   r6   rL   r(   r)   ra   r8   r9   r:   store	storedist	left_argsr7   rT   s                    r#   
_georadiuszGeoCommandsMixin._georadiusa   s    `lb %',	`
 `
 `
\QHhy$yS\ $47IINN'Q000'1#)S$PTV[]fhl'm'm 	$##E;%#HHH{### 	$##I{d#KKK{###k9h???r%   GEOADD)r+   fixedrepeatc           	         t          |ddd          \  \  }}}}|r|rt          t          j                  t	          |          dk    st	          |          dz  dk    rt          t          j                  |j        }t	          |          d}	}t          dt	          |          d          D ]}}
t          j	        ||
dz                      t          j	        ||
dz                      ||
dz            }}}||v r|r||vr,|s*|
                    |t          ||d	                    r|	dz  }	~|	r|                                 |r|	S t	          |          |z
  S )
N)nxxxchFTrd   r      r      
   )r	   r   r    NX_XX_GT_LT_ERROR_MSGr0   SYNTAX_ERROR_MSGrY   ranger   r   r[   r   updated)rQ   r@   rL   rr   rq   rs   datar5   old_lenchanged_itemsir.   r/   r+   s                 r#   geoaddzGeoCommandsMixin.geoaddv   sw   ) %'+	
 
 
Rd  	:" 	:d8999t99>>SYY]a//d3444y!$TAq#d))Q'' 	' 	'AT!a%[))T!a%[))QU #D
 RT-=-=b-=88D*S$";";<< '!Q&M 	KKMMM 	!  4yy7""r%   GEOHASHmembersc                 R    t          |j        j        |          }d |D             }|S )Nc                 D    g | ]}||dz                                    n|S )N0)r-   .0r>   s     r#   
<listcomp>z,GeoCommandsMixin.geohash.<locals>.<listcomp>   s3    1o1o1ocdAG3C3C3E3E3E\]1o1o1or%   maprY   r   )rQ   r@   r   hashesgeohash_lists        r#   geohashzGeoCommandsMixin.geohash   s/    SY]G,,1o1ohn1o1o1or%   GEOPOSc                 f    d t          |j        j        |          D             }d |D             }|S )Nc              3   <   K   | ]}|t          |          n|V  d S r=   r   r   s     r#   	<genexpr>z*GeoCommandsMixin.geopos.<locals>.<genexpr>   s0      cc!
1Accccccr%   c                 6    g | ]}||d         |d         gnd S )Nr   r    r   s     r#   r   z+GeoCommandsMixin.geopos.<locals>.<listcomp>   s,    OOOQ1qt4OOOr%   r   )rQ   r@   r   gospositionsr1   s        r#   geoposzGeoCommandsMixin.geopos   s<    ccs39=ZaGbGbcccOO,OOO
r%   GEODISTm1m2c                    |j                             |          |j                             |          g}t          d |D                       rd S d |D             }t          |d         d         |d         d         f|d         d         |d         d         f          }t	          |          dk    rt          |d                   nd}||z  S )Nc              3      K   | ]}|d u V  	d S r=   r   )r   elems     r#   r   z+GeoCommandsMixin.geodist.<locals>.<genexpr>   s&      22tt|222222r%   c                 ,    g | ]}t          |          S r   r   r   s     r#   r   z,GeoCommandsMixin.geodist.<locals>.<listcomp>   s    555aJqMM555r%   r   r   )rY   r   rc   r   r0   r$   )	rQ   r@   r   r   rL   	geohashesgeo_locsr1   r"   s	            r#   geodistzGeoCommandsMixin.geodist   s    Y]]2&&	b(9(9:	22	22222 	4559555AA7(1+a.(ST+VW.9YZZ47IINN'Q000Tzr%   GEORADIUS_ROc           
         t          |ddd          \  \  }}}}	}
}}|	pt          j        }	t          |          dk    rt	          |d                   nd}t          |j        |||||	|
|          }t          |||          }|S )N)r(   r)   ra   rb   rc   r:   Frd   r   r   g      ?)r	   rg   rh   r0   r$   rI   rY   r4   )rQ   r@   r.   r/   r6   rL   r(   r)   ra   r8   r9   r:   rk   r7   rT   rets                   r#   georadius_rozGeoCommandsMixin.georadius_ro   s    MYJ %',	N
 N
 N
J?Hhy$ $;>t99>>.tAw777s CvtUIW[\\[)X>>
r%   	GEORADIUSc                 $     | j         ||||g|R  S r=   )rl   )rQ   r@   r.   r/   r6   rL   s         r#   	georadiuszGeoCommandsMixin.georadius   s$     tsD#v=====r%   GEORADIUSBYMEMBERmember_namec                     |j                             |          }t          |          \  }}}} | j        ||||g|R  S r=   )rY   r   r   rl   	rQ   r@   r   r6   rL   member_scorer/   r.   rG   s	            r#   georadiusbymemberz"GeoCommandsMixin.georadiusbymember   sJ    y}}[11$\22T1atsD#v=====r%   GEORADIUSBYMEMBER_ROc                     |j                             |          }t          |          \  }}}} | j        ||||g|R  S r=   )rY   r   r   r   r   s	            r#   georadiusbymember_roz%GeoCommandsMixin.georadiusbymember_ro   sK    y}}[11$\22T1a t dC?$????r%   	GEOSEARCHc                     t          |ddd          \  \  }\  }}}}||t          t          j                  ||t          t          j                  |r | j        |||g|R  S  | j        ||||g|R  S )N)*frommember..fromlonlat	.byradiusFrd   )r	   r   r    rx   r   r   )rQ   r@   rL   
frommemberr.   r/   r6   rk   s           r#   	geosearchzGeoCommandsMixin.geosearch   s    7C8 %',	8
 8
 8
4)[dC&9 $,d3444!d&6d3444 	I,4,S*fQyQQQQ$4$S$VHiHHHHr%   GEOSEARCHSTOREdstsrcc                    t          |ddd          \  \  }\  }}}}}	||t          t          j                  ||t          t          j                  |rd|gnd|g}
|r | j        |||g|	|
R  S  | j        ||||g|	|
R  S )N)r   r   r   rj   Frd   s	   storedists   store)r	   r   r    rx   r   rl   )rQ   r   r   rL   r   r.   r/   r6   rj   rk   
additionals              r#   geosearchstorezGeoCommandsMixin.geosearchstore   s     COE %',	C
 C
 C
?4[dC&)i $,d3444!d&6d3444,5JlC((Hc?
 	T)4)#z6[I[PZ[[[["4?3c6SIS
SSSSr%   )__name__
__module____qualname__r   rP   bytesr   r&   boolintr_   r   floatr   rl   r
   r   r   r   r   r   r   r   r   r   r   r   r   r   r   __classcell__)rR   s   @r#   rK   rK   S   s5       c S T       E  Y  \`  eh        @@&+@27@AF@OT@	tE{C	 @ @ @ @* W(33t99,x@@@#+ #e # # # # A@#8 W)CCIIu#5uhGGG; % Dud{AS<T    HG
 W(33t99e"4eXFFF+  4e@U;V    GF
 W)CCIIue#<eXNNN; E u U xX]    ON W.TE5%(HRWQYZZZ 5 u e \a fjknfo    [Z W+cc$iiu%EuhWWW>>&+>27>AF>OT>	tCy#~	> > > XW>
 W%cc$ii-FPUxXXX>[ >u >e >\a >fjknfo > > > YX>
 W(TE50ISXRZ[[[@ @% @QV @_d @imnqir @ @ @ \[@
 W+cc$ii\5(CCCI[ I I49 I I I DCI  WCII
 x  T% Tk T% TDQTI T T T T T T T Tr%   rK   )%rg   collectionsr   typingr   r   r   r   	fakeredisr   r    fakeredis._command_args_parsingr	   fakeredis._commandsr
   r   r   r   fakeredis._helpersr   r   fakeredis.geor   r   r   fakeredis.modelr   r   r   r   r$   r&   r   r4   r   rI   rK   r   r%   r#   <module>r      s	   



 " " " " " " - - - - - - - - - - - - # # # # # # 8 8 8 8 8 8 @ @ @ @ @ @ @ @ @ @ @ @ 4 4 4 4 4 4 4 4 : : : : : : : : : :            7CC	u      J{$ABB	$y/ d d tTWy    & 
 	    	 
         
)_       F`T `T `T `T `T `T `T `T `T `Tr%   