B
    ŻĻfėE  ć               @  s  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$m(Z( d?d&d'd(d)d*d+d,Z)d&d'd-d.d/d0d1d2d3Z*d&d'd4d/d0d5d6d7Z+d&d'd8d/d0d5d9d:Z,d&d'd0d;d<d=Z-d>S )@é    )Śannotations)ŚList)ŚOptional)ŚUnion)ŚAssignmentStmt)ŚCallExpr)ŚClassDef)Ś	Decorator)Ś
LambdaExpr)ŚListExpr)Ś
MemberExpr)ŚNameExpr)ŚPlaceholderNode)ŚRefExpr)ŚStrExpr)Ś
SymbolNode)ŚSymbolTableNode)ŚTempNode)ŚTypeInfo)ŚVar)ŚSemanticAnalyzerPluginInterface)ŚAnyType)ŚCallableType)Śget_proper_type)ŚInstance)ŚNoneType)Ś
ProperType)ŚType)Ś	TypeOfAny)ŚUnboundType)Ś	UnionTypeé   )Śapply)Śinfer)Śnames)ŚutilFr   r   Śboolz(Optional[List[util.SQLAlchemyAttribute]])ŚclsŚapiŚis_mixin_scanŚreturnc             C  s  t  | |”}|d krd S | j d”r(d S t  ||”}t  |” |d k	r\|sXt | ||” |S g }| jj	sxr|j
 ” D ]\}}t| |||| qtW nJxHt  | jj	”D ]6}t|trĄt| ||| q¢t|tr¢t| ||| q¢W t| | |sųt | ||” t  ||” |S )NŚbuiltins)r%   Śinfo_for_clsŚfullnameŚ
startswithZget_mapped_attributesZestablish_as_sqlalchemyr"   Z re_apply_declarative_assignmentsŚdefsŚbodyr$   ŚitemsŚ_scan_symbol_table_entryZflatten_typecheckingŚ
isinstancer   Ś!_scan_declarative_assignment_stmtr	   Ś _scan_declarative_decorator_stmtŚ_scan_for_mapped_basesZadd_additional_orm_attributesZset_mapped_attributes)r'   r(   r)   ŚinfoZmapped_attributesZsym_nameŚsymŚstmt© r:   ś`C:\Users\Lenovo\AppData\Local\Temp\pip-unpacked-wheel-4nia5yor\sqlalchemy\ext\mypy\decl_class.pyŚ,scan_declarative_assignments_and_apply_types0   s8    



r<   Śstrr   zList[util.SQLAlchemyAttribute]ŚNone)r'   r(   ŚnameŚvalueŚ
attributesr*   c          	   C  s  t |j}t|tsdS d}t |j”}d}|tjtjtjtj	tj
tjhkrh|jrbt |jd }nd}n²|tjkr|jsd}nt |jd }	t|	tr|	j}	t|	ttfr| |	j|	”}
|
dk	rt|
jtrt |
jtj”rtt ||
jg ”t g}nt |d |
jj”|” |rFd}t || || j”| ” ttj}|dk	r|jdk	s`t | !tj"||jj#|jj$|| j%d” dS )zaExtract mapping information from a SymbolTableNode that's in the
    type.names dictionary.

    NFr   Tz4Column type should be a TypeEngine subclass not '{}'zŲCan't infer type from attribute {} on class {}. please specify a return type from this function that is one of: Mapped[<python type>], relationship[<target class>], Column[<TypeEngine>], MapperProperty[<python type>])r?   ŚlineŚcolumnŚtypr7   )&r   Śtyper3   r   r$   Śtype_id_for_named_nodeŚMAPPEDŚRELATIONSHIPŚCOMPOSITE_PROPERTYŚMAPPER_PROPERTYŚSYNONYM_PROPERTYŚCOLUMN_PROPERTYŚargsŚCOLUMNr   r   Ślookup_qualifiedr?   ŚnodeŚhas_base_type_idŚ
TYPEENGINEr    r#   Ś#extract_python_type_from_typeenginer   r%   ŚfailŚformatr-   r   r   Śspecial_formŚAssertionErrorŚappendŚSQLAlchemyAttributerB   rC   r7   )r'   r(   r?   r@   rA   Z
value_typeŚleft_hand_explicit_typeŚtype_idŚerrŚtypeengine_argr8   Śmsgr:   r:   r;   r2   r   sb    






r2   r	   )r'   r(   r9   rA   r*   c          	   C  sx  x4|j D ]&}t|tttfrt |”tjkrP qW dS | jj	 
|”}d}t |j”rttj}t|jj}|j|_t|gt|}	|jj|	_|	| jj	|< dS t|jjtr¤|jjj}
t|
trŠt |
| |”}ndS |tjtjtjtjtjtj hkr|
j!rt"|
j!d }n|tj#kr¤|
j!r¤|
j!d }t|tr¤| $|j|”}|dk	r¤t|jt%r¤t &|jtj'”rt(t) *||jg ”t+ g}nt ,|d -|jj.”|
” |dkrŌd}t ,|| -|jj”|” ttj}t|jj}|j|_t|trt"t /||”}| 0tj1|g”|j_t 2t3|jj4|jj	”}t|g|}	|jj|	_| 5tj6|j|j7|j8|| j9d” |	| jj	|< dS )a°  Extract mapping information from a @declared_attr in a declarative
    class.

    E.g.::

        @reg.mapped
        class MyClass:
            # ...

            @declared_attr
            def updated_at(cls) -> Column[DateTime]:
                return Column(DateTime)

    Will resolve in mypy as::

        @reg.mapped
        class MyClass:
            # ...

            updated_at: Mapped[Optional[datetime.datetime]]

    Nr   z4Column type should be a TypeEngine subclass not '{}'ząCan't infer type from @declared_attr on function '{}';  please specify a return type from this function that is one of: Mapped[<python type>], relationship[<target class>], Column[<TypeEngine>], MapperProperty[<python type>])r?   rB   rC   rD   r7   ):Z
decoratorsr3   r   r   r   r$   rF   ZDECLARED_ATTRr/   r0   Śindexr%   Zname_is_dunderr?   r   r   rV   ŚvarrP   r   r   rE   Śfuncr   Zret_typer   Ztype_id_for_unbound_typerG   rH   rI   rJ   rK   rL   rM   r   rN   rO   r   rQ   rR   r    r#   rS   r   rT   rU   r-   Zunbound_to_instanceZ
named_typeZNAMED_TYPE_SQLA_MAPPEDZexpr_to_mapped_constructorr
   Ś	argumentsrX   rY   rB   rC   r7   )r'   r(   r9   rA   ŚdecZ	dec_indexrZ   Zany_Z	left_nodeZnew_stmtZ	func_typer[   r]   r8   r^   Śrvaluer:   r:   r;   r5   Ė   s    









r5   r   c          	   C  s  |j d }t|tsdS | jj |j”}|dk	s4t|j}t|t	rHdS ||jksVtt|t
sdt|jdkr| |j”dkrt | j” dS |jdkr¦t | j” nh|j d”r¶dS |jdkrt|jtsŽt |d|” n0x.|jjD ]"}t|ttfrčt | |||” qčW d}d}	|js*|jdkrt|jtrą|j}	|jjd	krą| d	| ”}
|
dk	rą|
jdk	rąt |
j”tjkrąt|jjd }	|j}nFt|j}t|trŲt |j”tjkrŲt|jd }	|}n|}	d}t|jt rž|dk	rž|	}nHt|jt!rBt|jj"t#rBt$ %||||	|jj"”}|dkrFdS ndS |dk	sTt| &tj'|j|j(|j)|| jd
” t *||||	|” dS )zZExtract mapping information from an assignment statement in a
    declarative class.

    r   NZ__abstract__TZ__tablename__Ś__Z_mypy_mapped_attrsz+_mypy_mapped_attrs is expected to be a listZMapped)r?   rB   rC   rD   r7   )+Zlvaluesr3   r   r7   r$   Śgetr?   rW   rP   r   r   Z
parse_boolrd   r%   Zset_is_baseZset_has_tabler.   r   rT   r1   r   r"   Zapply_mypy_mapped_attrZis_inferredrE   r   rO   rF   rG   r   rM   r   r   r   Zcalleer   r#   Z#infer_type_from_right_hand_nameexprrX   rY   rB   rC   Zapply_type_to_mapped_statement)r'   r(   r9   rA   Zlvaluer8   rP   ŚitemZleft_hand_mapped_typerZ   Z
mapped_symZ	node_typeZpython_type_for_typer:   r:   r;   r4   g  s    










r4   )r'   r(   r*   c             C  sR   t  | |”}|dkrdS x4|jdd D ]"}|j d”r:q(t|j|dd q(W dS )zčGiven a class, iterate through its superclass hierarchy to find
    all other classes that are considered as ORM-significant.

    Locates non-mapped mixins and scans them for mapped attributes to be
    applied to subclasses.

    Nr!   é’’’’r+   T)r)   )r%   r,   Śmror-   r.   r<   Zdefn)r'   r(   r7   Z	base_infor:   r:   r;   r6   é  s    r6   N)F).Ś
__future__r   Śtypingr   r   r   Z
mypy.nodesr   r   r   r	   r
   r   r   r   r   r   r   r   r   r   r   r   Zmypy.pluginr   Z
mypy.typesr   r   r   r   r   r   r   r   r   r    Ś r"   r#   r$   r%   r<   r2   r5   r4   r6   r:   r:   r:   r;   Ś<module>   sR   ?Y  