U
    cc2                   	   @   s  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( ed#eeef d$Z)G d%d& d&Z*d'd( Z+ee,edd)d*d+Z-ee,ee d,d-d.Z.ee,ee d,d/d0Z/edd1d2d3Z0edd1d4d5Z1ee2d1d6d7Z3ee2d1d8d9Z4edd1d:d;Z5ee2d1d<d=Z6eeeee*  d>d?d@Z7eee* ddAdBdCZ8ee,eddDdEdFZ9eeef e,e,e,ddGdHdIZ:edJdKee,deeeef  dLdMdNZ;eee,ee
e) dJf ee) dLdOdNZ;ddKee,eee
e dJf  ee dLdPdNZ;ee ee dQdRdSZ<ee	e	dTdUdVZ=eeee dWdXdYZ>eedZd[d\Z?dS )]    N)Any)Iterable)Iterator)List)Optional)overload)Tuple)Type)TypeVar)Union)ARG_POS)CallExpr)ClassDef)CLASSDEF_NO_INFO)Context)
Expression)IfStmt)JsonDict)
MemberExpr)NameExpr)	Statement)SymbolTableNode)TypeInfo)ClassDefContext)DynamicClassDefContext)SemanticAnalyzerPluginInterface)deserialize_and_fixup_type)map_type_from_supertype)Instance)NoneType)TypeVarType)UnboundType)	UnionType	_TArgType)boundc                   @   s^   e Zd Zeeeee eddddZe	dddZ
eddd	d
Zeee	ed dddZdS )SQLAlchemyAttributeN)namelinecolumntypinforeturnc                 C   s"   || _ || _|| _|| _|| _d S N)r&   r'   r(   typer*   )selfr&   r'   r(   r)   r*    r/   </tmp/pip-unpacked-wheel-8u86ls_i/sqlalchemy/ext/mypy/util.py__init__+   s
    zSQLAlchemyAttribute.__init__)r+   c                 C   s$   | j s
t| j| j| j| j  dS )N)r&   r'   r(   r-   )r-   AssertionErrorr&   r'   r(   	serialize)r.   r/   r/   r0   r3   9   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*   )r.   r4   r/   r/   r0   expand_typevar_from_subtypeB   s    z/SQLAlchemyAttribute.expand_typevar_from_subtype)r*   dataapir+   c                 C   s,   |  }t|d|}| f ||d|S )Nr-   )r)   r*   )copyr   pop)clsr*   r7   r8   r)   r/   r/   r0   deserializeK   s    zSQLAlchemyAttribute.deserialize)__name__
__module____qualname__strintr   r	   r   r1   r   r3   r6   classmethodr   r<   r/   r/   r/   r0   r%   *   s   		r%   c                 C   s   t td| S )Nz	^__.+?__$)boolrematch)r&   r/   r/   r0   name_is_dunderW   s    rF   )r*   keyr7   r+   c                 C   s   || j di |< d S NZ
sqlalchemymetadata
setdefault)r*   rG   r7   r/   r/   r0   _set_info_metadata[   s    rL   )r*   rG   r+   c                 C   s   | j di |d S rH   )rJ   get)r*   rG   r/   r/   r0   _get_info_metadata_   s    rN   c                 C   s0   | j r,| j D ]}t||}|d k	r|  S qd S r,   )mrorN   )r*   rG   baserJ   r/   r/   r0   _get_info_mro_metadatac   s    


rQ   )r*   r+   c                 C   s   | j di  d S rH   rI   r*   r/   r/   r0   establish_as_sqlalchemyl   s    rS   c                 C   s   t | dd d S Nis_baseTrL   rR   r/   r/   r0   set_is_basep   s    rW   c                 C   s   t | d}|dkS rT   rN   r*   rU   r/   r/   r0   get_is_baset   s    
rZ   c                 C   s   t | d}|dkS rT   )rQ   rY   r/   r/   r0   has_declarative_basey   s    
r[   c                 C   s   t | dd d S NZ	has_tableTrV   rR   r/   r/   r0   set_has_table~   s    r]   c                 C   s   t | d}|dkS r\   rX   rY   r/   r/   r0   get_has_table   s    
r^   )r*   r8   r+   c                 C   sJ   t | d}|d krd S g }|D ]&}t| ||}||  || q|S )Nmapped_attributes)rN   r%   r<   r6   append)r*   r8   r_   
attributesr7   attrr/   r/   r0   get_mapped_attributes   s     
rc   )r*   ra   r+   c                 C   s   t | ddd |D  d S )Nr_   c                 S   s   g | ]}|  qS r/   )r3   ).0	attributer/   r/   r0   
<listcomp>   s     z)set_mapped_attributes.<locals>.<listcomp>rV   )r*   ra   r/   r/   r0   set_mapped_attributes   s
    rg   )r8   msgctxr+   c                 C   s   d| }|  ||S )Nz[SQLAlchemy Mypy plugin] %s)fail)r8   rh   ri   r/   r/   r0   rj      s    rj   )ri   modulesymbol_nameasnamer+   c                 C   s8   | j j| j j j}||kr4| j j| j| }|||< d S r,   )r8   modulesZ
cur_mod_idnames)ri   rk   rl   rm   module_globalsZ
lookup_symr/   r/   r0   
add_global   s    rq   .)
expr_types)callexprr&   rr   r+   c                C   s   d S r,   r/   rs   r&   rr   r/   r/   r0   get_callexpr_kwarg   s    ru   c                C   s   d S r,   r/   rt   r/   r/   r0   ru      s    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,   )	arg_namesindex
ValueErrorargsr5   r   r   )rs   r&   rr   Zarg_idxkwargr/   r/   r0   ru      s    
 )stmtsr+   c                 c   sX   | D ]N}t |trLt |jd trL|jd jdkrL|jd jD ]
}|V  q>q|V  qd S )Nr   ztyping.TYPE_CHECKING)r5   r   exprr   fullnamebody)r{   ZstmtZsubstmtr/   r/   r0   flatten_typechecking   s    
r   )r8   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)rd   Ztyp_argr8   r/   r0   rf      s     z'unbound_to_instance.<locals>.<listcomp>Nc                    s$   g | ]}t |trt |n|qS r/   )r5   r!   r   )rd   argr   r/   r0   rf     s   )r5   r!   r&   r   r"   ry   r   lookup_qualifiedr   noder   r   )r8   r)   r   Z
bound_typer/   r   r0   r      s4    	




r   )r;   r8   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   r&   r5   r   r   r2   )r;   r8   symr/   r/   r0   info_for_cls  s    
r   )r|   r+   c                 C   s,   t d}d|_t|d}t|| gtgdgS )NZ__sa_Mappedz sqlalchemy.orm.attributes.MappedZ_empty_constructorZarg1)r   r}   r   r   r   )r|   Zcolumn_descriptorZmember_exprr/   r/   r0   expr_to_mapped_constructor(  s    
r   )@rD   typingr   r   r   r   r   r   r   r	   Z
TypingTyper
   r   Z
mypy.nodesr   r   r   r   r   r   r   r   r   r   r   r   r   Zmypy.pluginr   r   r   Zmypy.plugins.commonr   Zmypy.typeopsr   Z
mypy.typesr   r   r    r!   r"   r#   r%   rF   r@   rL   rN   rQ   rS   rW   rC   rZ   r[   r]   r^   rc   rg   rj   rq   ru   r   r   r   r   r/   r/   r/   r0   <module>   s   -	 
 

   1 