o
    h#                     @  s   U d dl mZ d dlZd dlZd dlmZmZ d dlmZ dZ	de
d< d,ddZ			d-d.ddZd/ddZdd Zdd Zdd  Zd0d!d"Zd0d#d$Zd1d%d&Zd2d(d)Zd3d*d+ZdS )4    )annotationsN)AnyFinal)config__streamlitmagic__r   MAGIC_MODULE_NAMEcodestrscript_pathreturnr   c                 C  s*   t | |d}t|| }t|d|d |S )a  Modifies the code to support magic Streamlit commands.

    Parameters
    ----------
    code : str
        The Python code.
    script_path : str
        The path to the script file.

    Returns
    -------
    ast.Module
        The syntax tree for the code.

    execT)is_rootfile_ends_in_semicolon)astparse_does_file_end_in_semicolon_modify_ast_subtree)r   r
   treer    r   X/var/www/vscode/kcb/lib/python3.10/site-packages/streamlit/runtime/scriptrunner/magic.py	add_magic   s   
r   bodyFr   	body_attrr   boolr   c              	   C  s  t | |}t|D ]\}}t|}|tju s%|tju s%|tju s%|tju r*t| q	|tj	u s9|tj
u s9|tju rCt| t|d q	|tju r]|jD ]}t|tjtjhv r[t| qKq	|tju sltjdkr|tju rt| t|dd t|dd |jD ]}	t|	 qq	|tju rt| t|d q	tjdkr|tju r|jD ]}
t|
 qq	|tju rt||t| ||t|d k|d}|dur||_q	|rt|  t|  dS )	z;Parses magic commands and modifies the given AST (sub)tree.orelse)      	finalbody)r   )r   
      )parent_typer   is_last_exprr   N)getattr	enumeratetyper   FunctionDefWithAsyncFunctionDef	AsyncWithr   ForWhileAsyncForClassDefr   Trysysversion_infoTryStarhandlersIfMatchcasesExpr_get_st_write_from_exprlenvalue_insert_import_statementfix_missing_locations)r   r   r   r   r   inode	node_type
inner_nodehandler_node	case_noder8   r   r   r   r   9   s^   














r   Nonec                 C  s   t  }| jrt| jd tjtjhv r| jd| dS t| jdkrKt| jd tju rKt	| jd j
rKt| jd tjtjhv rK| jd| dS | jd| dS )z>Insert Streamlit import statement at the top(ish) of the tree.r   r      N)_build_st_import_statementr   r$   r   
ImportFromImportinsertr7   r5   _is_string_constant_noder8   )r   	st_importr   r   r   r9      s    r9   c                   C  s   t jt jdtdgdS )z>Build AST node for `import magic_funcs as __streamlitmagic__`.z*streamlit.runtime.scriptrunner.magic_funcs)nameasname)names)r   rE   aliasr   r   r   r   r   rC      s   rC   c                 C  s.   t jt jdt jtt  dt  d| g dS )zBBuild AST node for `__streamlitmagic__.transparent_write(*nodes)`.transparent_write)idctx)attrr8   rO   )funcargskeywords)r   Call	AttributeNamer   Load)nodesr   r   r   _build_st_write_call   s   rY   c                 C  s   t | jtju rt|||sd S t| j||rt|sd S t | jtju s-t | jtju r/d S t | jtj	u r9d S t | jtj
u rE| jjn| jg}t|S N)r$   r8   r   rT   _is_displayable_last_expr_is_docstring_node*_should_display_docstring_like_node_anywayYield	YieldFromAwaitTupleeltsrY   )r<   r;   r    r   r!   r   rR   r   r   r   r6      s      r6   c                 C  s   t | tjot | jtS rZ   )
isinstancer   Constantr8   r	   )r<   r   r   r   rG      s   rG   c                 C  s$   |dkot | o|tjtjtjhv S )Nr   )rG   r   r%   r'   Module)r<   
node_indexr    r   r   r   r\      s
   r\   c                 C  s`   d}t dr.t| jdkrdS t| jd dd }|d ur.|d|d  }|dd	}|S )
NF"magic.displayLastExprIfNoSemicolonr   
end_lineno
r    ;)r   
get_optionr7   r   r"   splitstripendswith)r   r   r   last_line_numlast_line_strr   r   r   r      s   
r   r!   c                 C  s   |o| o| ot dS )Nrg   r   rm   )r   r!   r   r   r   r   r[      s   r[   c                 C  s   t do| S )Nzmagic.displayRootDocStringrs   )r   r   r   r   r]     s   r]   )r   r	   r
   r	   r   r   )r   FF)r   r   r   r	   r   r   r   r   )r   r   r   rA   )r   r   )r   r	   r   r   )r   r   r!   r   r   r   r   r   )r   r   r   r   )
__future__r   r   r.   typingr   r   	streamlitr   r   __annotations__r   r   r9   rC   rY   r6   rG   r\   r   r[   r]   r   r   r   r   <module>   s(   

Q
$


