U
    [h&                     @  s  d dl mZ d dlZd dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlm	Z	 d d	lm
Z
 d d
lmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlm Z  d dlm!Z! d dlm"Z" d dl#m$Z$ d dl%m&Z& d dl%m'Z' d d l%m(Z( d d!l)m*Z* d d"l+m,Z, d d#l-m.Z. d d$l-m/Z/ d d%l-m0Z0 d d&l-m1Z1 d d
l-mZ d d'l-m2Z2 d d(l-m3Z3 d d)l-m4Z4 e5d*d+ ej67d,D Z8e8d-kZ9ed.eeef d/Z:G d0d1 d1Z;d2d3d4d5d6Z<d7d2d8d9d:d;d<Z=d7d2d=d>d?d@Z>d7d2d=d>dAdBZ?d7d9dCdDdEZ@d7d9dCdFdGZAd7d3dCdHdIZBd7d3dCdJdKZCd7d9dCdLdMZDd7d3dCdNdOZEd7dPdQdRdSdTZFdUdVd2dWdXdYZd7dZd9d[d\d]ZGdPd2d^d9d_d`daZHdbd2d2d2d9dcdddeZIe	dfdgdhd2d9didjdkdlZJe	dhd2dmdndjdodlZJddgdhd2dpd=djdqdlZJdrdsdtdudvZKdwdxdydzd{ZLdPdUdUd|d}d~ZMddPddddZNdUddddZOdS )    )annotationsN)Any)Iterable)Iterator)List)Optional)overload)Tuple)Type)TypeVar)Union)version)format_type)CallExpr)ClassDef)CLASSDEF_NO_INFO)Context)
Expression)FuncDef)IfStmt)JsonDict)
MemberExpr)NameExpr)	Statement)SymbolTableNode)	TypeAlias)TypeInfo)Options)ClassDefContext)DynamicClassDefContext)SemanticAnalyzerPluginInterface)deserialize_and_fixup_type)map_type_from_supertype)CallableType)get_proper_type)Instance)NoneType)TypeVarType)UnboundType)	UnionTypec                 C  s    g | ]}t d |rt|qS )z^\d+$)rematchint).0x r/   >./venv/lib/python3.8/site-packages/sqlalchemy/ext/mypy/util.py
<listcomp>6   s      r1   .)      	_TArgType)Zboundc                   @  sZ   e Zd ZdddddddddZd	d
ddZdddddZedd	dd dddZdS )SQLAlchemyAttributestrr,   zOptional[Type]r   None)namelinecolumntypinforeturnc                 C  s"   || _ || _|| _|| _|| _d S N)r9   r:   r;   typer=   )selfr9   r:   r;   r<   r=   r/   r/   r0   __init__?   s
    zSQLAlchemyAttribute.__init__r   )r>   c                 C  s$   | j s
t| j| j| jt| j dS )N)r9   r:   r;   r@   )r@   AssertionErrorr9   r:   r;   serialize_type)rA   r/   r/   r0   	serializeM   s    
zSQLAlchemyAttribute.serialize)sub_typer>   c                 C  s&   t | jtsdS t| j|| j| _dS )zxExpands type vars in the context of a subtype when an attribute is
        inherited from a generic super type.
        N)
isinstancer@   r'   r"   r=   )rA   rF   r/   r/   r0   expand_typevar_from_subtypeV   s    z/SQLAlchemyAttribute.expand_typevar_from_subtyper    )r=   dataapir>   c                 C  s,   |  }t|d|}| f ||d|S )Nr@   )r<   r=   )copyr!   pop)clsr=   rI   rJ   r<   r/   r/   r0   deserialize_   s    zSQLAlchemyAttribute.deserializeN)__name__
__module____qualname__rB   rE   rH   classmethodrN   r/   r/   r/   r0   r6   >   s
   		r6   r7   bool)r9   r>   c                 C  s   t td| S )Nz	^__.+?__$)rS   r*   r+   )r9   r/   r/   r0   name_is_dunderk   s    rT   r   r   r8   )r=   keyrI   r>   c                 C  s   || j di |< d S NZ
sqlalchemymetadata
setdefault)r=   rU   rI   r/   r/   r0   _set_info_metadatao   s    rZ   zOptional[Any])r=   rU   r>   c                 C  s   | j di |d S rV   )rX   get)r=   rU   r/   r/   r0   _get_info_metadatas   s    r\   c                 C  s0   | j r,| j D ]}t||}|d k	r|  S qd S r?   )mror\   )r=   rU   baserX   r/   r/   r0   _get_info_mro_metadataw   s    


r_   )r=   r>   c                 C  s   | j di  d S rV   rW   r=   r/   r/   r0   establish_as_sqlalchemy   s    ra   c                 C  s   t | dd d S Nis_baseTrZ   r`   r/   r/   r0   set_is_base   s    re   c                 C  s   t | d}|dkS rb   r\   r=   rc   r/   r/   r0   get_is_base   s    
rh   c                 C  s   t | d}|dkS rb   )r_   rg   r/   r/   r0   has_declarative_base   s    
ri   c                 C  s   t | dd d S NZ	has_tableTrd   r`   r/   r/   r0   set_has_table   s    rk   c                 C  s   t | d}|dkS rj   rf   rg   r/   r/   r0   get_has_table   s    
rl   r    z#Optional[List[SQLAlchemyAttribute]])r=   rJ   r>   c                 C  sJ   t | d}|d krd S g }|D ]&}t| ||}||  || q|S )Nmapped_attributes)r\   r6   rN   rH   append)r=   rJ   rm   
attributesrI   attrr/   r/   r0   get_mapped_attributes   s     
rq   r
   r   )typ_optionsr>   c                 C  s   t rt| |S t| S d S r?   )mypy_14_mypy_format_type)rr   rs   r/   r/   r0   r      s    
r   zList[SQLAlchemyAttribute])r=   ro   r>   c                 C  s   t | ddd |D  d S )Nrm   c                 S  s   g | ]}|  qS r/   )rE   )r-   Z	attributer/   r/   r0   r1      s     z)set_mapped_attributes.<locals>.<listcomp>rd   )r=   ro   r/   r/   r0   set_mapped_attributes   s
    rv   r   )rJ   msgctxr>   c                 C  s   d| }|  ||S )Nz[SQLAlchemy Mypy plugin] %s)fail)rJ   rw   rx   r/   r/   r0   ry      s    ry   z.Union[ClassDefContext, DynamicClassDefContext])rx   modulesymbol_nameasnamer>   c                 C  s8   | j j| j j j}||kr4| j j| j| }|||< d S r?   )rJ   modulesZ
cur_mod_idnames)rx   rz   r{   r|   module_globalsZ
lookup_symr/   r/   r0   
add_global   s    r   .)
expr_typesr   z#Optional[Union[CallExpr, NameExpr]])callexprr9   r   r>   c                C  s   d S r?   r/   r   r9   r   r/   r/   r0   get_callexpr_kwarg   s    r   z!Tuple[TypingType[_TArgType], ...]zOptional[_TArgType]c                C  s   d S r?   r/   r   r/   r/   r0   r      s    z%Optional[Tuple[TypingType[Any], ...]]c                C  sT   z| j |}W n tk
r&   Y d S X | j| }t||d k	rB|nttfrP|S d S r?   )Z	arg_namesindex
ValueErrorargsrG   r   r   )r   r9   r   Zarg_idxZkwargr/   r/   r0   r      s    
 zIterable[Statement]zIterator[Statement])stmtsr>   c                 c  sR   | D ]H}t |trFt |jd trF|jd jdkrF|jd jE d H  q|V  qd S )Nr   ztyping.TYPE_CHECKING)rG   r   exprr   fullnameZbody)r   Zstmtr/   r/   r0   flatten_typechecking   s    r   r   z#Optional[Union[Instance, TypeInfo]])calleer>   c                 C  s   t | ttfrt | jtrP| jjrLt | jjtrLt| jjj}t |t	rL|S d S t | jt
rxt| jj}t |t	r|S nt | jtr| jS d S r?   )rG   r   r   noder   r@   r#   r$   ret_typer%   r   targetr   )r   r   Ztarget_typer/   r/   r0   type_for_callee  s    

r   )rJ   r<   r>   c                   s   t |ts|S |jdkr>t t fdd|jD t g S  |j|}|dk	rt |trt |j	t
r|j	}t| fdd|jD S |S dS )zTake the UnboundType that we seem to get as the ret_type from a FuncDef
    and convert it into an Instance/TypeInfo kind of structure that seems
    to work as the left-hand type of an AssignmentStatement.

    r   c                   s   g | ]}t  |qS r/   )unbound_to_instance)r-   Ztyp_argrJ   r/   r0   r1   +  s     z'unbound_to_instance.<locals>.<listcomp>Nc                   s$   g | ]}t |trt |n|qS r/   )rG   r(   r   )r-   argr   r/   r0   r1   ;  s   )rG   r(   r9   r   r)   r   r&   lookup_qualifiedr   r   r   r%   )rJ   r<   r   Z
bound_typer/   r   r0   r     s4    	



r   r   zOptional[TypeInfo])rM   rJ   r>   c                 C  sD   | j tkr>|| j| }|d kr$d S |r4t|jts8t|jS | j S r?   )r=   r   r   r9   rG   r   r   rC   )rM   rJ   Zsymr/   r/   r0   info_for_clsH  s    
r   zUnion[str, JsonDict])r<   r>   c                 C  s\   z
|   W S  tk
r   Y nX t| drBtdd | jD | _nt| drT|  } |   S )Nr   c                 s  s$   | ]}t |d r| n|V  qdS )resolve_string_annotationN)hasattrr   )r-   ar/   r/   r0   	<genexpr>[  s   
z!serialize_type.<locals>.<genexpr>r   )rE   	Exceptionr   tupler   r   )r<   r/   r/   r0   rD   U  s    


rD   )PZ
__future__r   r*   typingr   r   r   r   r   r   r	   r
   Z
TypingTyper   r   Zmypyr   Zmypy.messagesr   ru   Z
mypy.nodesr   r   r   r   r   r   r   r   r   r   r   r   r   r   Zmypy.optionsr   Zmypy.pluginr   r   r    Zmypy.plugins.commonr!   Zmypy.typeopsr"   Z
mypy.typesr#   r$   r%   r&   r'   r(   r)   r   __version__splitZ_versrt   r5   r6   rT   rZ   r\   r_   ra   re   rh   ri   rk   rl   rq   rv   ry   r   r   r   r   r   r   rD   r/   r/   r/   r0   <module>   s   -	
2