
    yj5                     $   d Z ddlZddlZddlZddlZddlZddlZddlZddlmZm	Z	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mZ dd
lmZmZmZmZ ddlmZ ddl m!Z!  G d de"          Z#d Z$d0dZ%d Z&d Z'd Z(d Z)d Z*d Z+d Z,d Z-d Z.d Z/d Z0eddfde1de2de3fd Z4d! Z5d"ee3ee3         f         fd#Z6d1d%Z7d& Z8d' Z9d(ee3e!e:ef         d)e1fd*Z;	 	 	 d2d(ee3e!e:ef         d.ee1         dee3         dee2         d)e"f
d/Z<dS )3uj   
本文件实现了常用的工具函数

Authors: xiangyiqing(xiangyiqing@baidu.com)
Date:    2023/07/24
    N)datetimetimezone	timedelta)log)FileIntegrityError)DEFAULT_MAX_WORKERSwrapstqdm)ThreadPoolExecutoras_completed)ListUnionBinaryIOOptional)VERSION)Pathc                       e Zd ZdZd Zd ZdS )Dictz
dict classc                 z    |                      |d           }t          |t                    rt          |          n|S N)get
isinstancedictr   selfkeyvalues      g/lsinfo/ai/hellotax_ai/data_center/backend/venv/lib/python3.11/site-packages/aistudio_sdk/utils/util.py__getattr__zDict.__getattr__$   s4    d##(55@tE{{{5@    c                     || |<   d S r    r   s      r    __setattr__zDict.__setattr__(   s    S			r"   N)__name__
__module____qualname____doc__r!   r%   r$   r"   r    r   r   "   s;        A A A    r"   r   c                 N    t          | t                    rt          |           S | S )z`
    Params
        :resp: dict, response from AIStudio
    Rerurns
        AIStudio object
    )r   r   r   )resps    r    convert_to_dict_objectr,   ,   s'     $ DzzKr"   c                     | |||dS )u  
    构造错误响应信息。

    Params:
        sdk_code (str): SDK错误码，标识错误类型。
        msg (str): 错误描述信息。
        biz_code (str, optional): 业务层面的错误码，透传自上游接口。
        log_id (str, optional): 与错误相关的日志ID，透传自上游接口。

    Returns:
        dict: 格式化好的错误信息。
    )
error_code	error_msgbiz_codelog_idr$   )sdk_codemsgr0   r1   s       r    err_respr4   9   s     	  r"   c                 h    t          j        dd| t           j                  } t          |           }|S )u   检测host合法性z
^https?:// )flags)resub
IGNORECASEis_valid_domain)hostresults     r    is_valid_hostr>   N   s2     6-T???DT""F Mr"   c                     dS )u   检测域名合法性Tr$   )domains    r    r;   r;   Y   s    4r"   c                     t          j                    }t          | d          5 }	 |                    d          }|sn|                    |           .	 ddd           n# 1 swxY w Y   |                                S )u   将文件计算为sha256值rbT   Nhashlibsha256openreadupdate	hexdigest)	file_pathsha256_hashfiledatas       r    calculate_sha256rO   `   s    .""K	i		 %$	%99U##D t$$$		% % % % % % % % % % % % % % % %   """s   0A!!A%(A%c                      t          t          d                    } t          j        |           }|                    d          }|S )uT   
    # 生成 ISO 8601日期时间格式
    # 例如"2023-09-12T11:29:45.703Z"
       )hoursmilliseconds)timespec)r   r   r   now	isoformat)zonerU   formatted_dates      r    gen_ISO_format_datestrrY   n   sD     IA&&&''D
,t

C]]N];;Nr"   c                 P   t          j                    }	 t          | d          5 }	 |                    d          }|sn|                    |           .	 ddd           n# 1 swxY w Y   n$# t
          $ r t          d|  d           Y dS w xY w|                                }|S )u   将文件计算为md5值rB   Ti   Nz
The file 'z' does not exist.)rE   md5rG   rH   rI   FileNotFoundErrorprintrJ   )rK   md5_hashrM   rN   md5_hexs        r    gen_MD5r`   {   s   {}}H
)T"" 	&d&yy %%%	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	& 	&    79777888tt
   ""GNs4   A. 0A"A. "A&&A. )A&*A. .BBc                 ~    |                      d          }t          j        |                              d          }|S )u   将字符串计算为base64utf-8)encodebase64	b64encodedecode)original_string
bytes_database64_encodeds      r    
gen_base64rj      s;     !''00J%j1188AANr"   c                 r   d|  d| }t          j        ddd          5 }|                    |           |j        }t	          j        |           ddd           n# 1 swxY w Y   	 t          |d          5 }t          j        |	                                          
                    d	          }ddd           n# 1 swxY w Y   |t          j                            |          rt          j        |           S S # t          j                            |          rt          j        |           w w xY w)
u?   生成指定内容的文件并进行base64编码字符串返回z6version https://git-lfs.github.com/spec/v1
oid sha256:z
size Fwz.txt)deletemodesuffixNrB   rb   )tempfileNamedTemporaryFilewritenamer   debugrG   rd   re   rH   rf   ospathexistsremove)rF   sizecontenttmptmp_pathfencodeds          r    $create_sha256_file_and_encode_base64r      s   ]]]W[]]G		$EF	K	K	K s		'8	(              
 (D!! 	AQ&qvvxx0077@@G	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A 	A7>>(## 	 Ih	 27>>(## 	 Ih	 sA   1AA!$A!)D  9:B?3D  ?CD  CD   6D6c                     t          | d          5 }|                                }t          j        |          }ddd           n# 1 swxY w Y   |                    d          S )u'   读取文件内容并进行Base64编码rB   Nrb   )rG   rH   rd   re   rf   )filenamerM   contentsencoded_contentss       r    file_to_base64r      s    	h		 699;;!+H556 6 6 6 6 6 6 6 6 6 6 6 6 6 6 ""7+++s   )AA
A
c           	      ~   t          j        |dt           j                  }t          j        |           D ]s\  }}}|                    | d          }|D ]T}|                    t          j                            ||          t          j                            ||                     Ut|	                                 dS )u   
    压缩指定文件夹
    :param dirpath: 目标文件夹路径
    :param out_full_name: 压缩文件保存路径 xxxx.zip
    :return: 无
    rl   r6   N)
zipfileZipFileZIP_DEFLATEDru   walkreplacerr   rv   joinclose)dirpathout_full_namezip_objrv   dirnames	filenamesfpathr   s           r    zip_dirr      s     omS'2FGGG%'WW%5%5 W W!h	Wb))! 	W 	WHMM"',,tX66UH8U8UVVVV	WMMOOOOOr"   c                    d}t          j                    }t          | d          5 }	 |                    |          }|sn|                    |           .	 ddd           n# 1 swxY w Y   |                                S )z
    hash
    rC   rB   TNrD   )rK   BUFFER_SIZErL   r}   rN   s        r    compute_hashr      s     K.""K	i		 %!	%66+&&D t$$$		% 	% % % % % % % % % % % % % % %   """s   0A##A'*A'c                     t          |           }||k    sCt          j        |            d| d|d|d}t          j        |           t          |          dS )zValidate the file hash is expected, if not, delete the file

    Args:
        file_path (str): The file to validate
        expected_sha256 (str): The expected sha256 hash

    Raises:
        FileIntegrityError: If file_path hash is not expected.

    zFile z6 integrity check failed, expected sha256 signature is z, actual is z3, the download may be incomplete, please try again.N)r   ru   rx   r   errorr   )rK   expected_sha256file_sha256r3   s       r    file_integrity_validationr      sp     y))K/))
	) II5 		# %%% *)r"   Fmax_workersdisable_tqdm	tqdm_descc                       fd}|S )af  
    A decorator to execute a function in a threaded manner using ThreadPoolExecutor.

    Args:
        max_workers (int): The maximum number of threads to use.
        disable_tqdm (bool): disable progress bar.
        tqdm_desc (str): Desc of tqdm.

    Returns:
        function: A wrapped function that executes with threading and a progress bar.


    c                 D     t                      fd            }|S )Nc           
         g }t          dddt          |           pdt          |            d          5 }t          
          5 	fd| D             }t          |          D ]>}|                    d	           |                    |                                           ?	 d d d            n# 1 swxY w Y   d d d            n# 1 swxY w Y   |S )
NT   r   zProcessing z items)
unit_scaleunit_divisorinitialtotaldescdisable)r   c                 6    i | ]} j         |gR i |S r$   )submit).0itemargsexecutorfunckwargss     r    
<dictcomp>zGthread_executor.<locals>.decorator.<locals>.wrapper.<locals>.<dictcomp>  sI         (dDTDDDVDDd  r"      )r   lenr   r   rI   appendr=   )iterabler   r   resultspbarfuturesfuturer   r   r   r   r   s    ``    @r    wrapperz3thread_executor.<locals>.decorator.<locals>.wrapper  s   G#!%h--"I&ICMM&I&I&I(   8 'K@@@ 
8H      $,  G #/w"7"7 8 8Av}}77778
8 
8 
8 
8 
8 
8 
8 
8 
8 
8 
8 
8 
8 
8 
88 8 8 8 8 8 8 8 8 8 8 8 8 8 8( Ns6   CA B:.C:B>	>CB>	CCCr	   )r   r   r   r   r   s   ` r    	decoratorz"thread_executor.<locals>.decorator   sC    	t	 	 	 	 	 	 	 
	2 r"   r$   )r   r   r   r   s   ``` r    thread_executorr      s0    "      < r"   c                    t           j                            d          r|                     dd          }n|                     dd          }t          j                            |d         |                    dd                    }|S )u   
    目录
    win\   /r   .___)sysplatform
startswithrsplitru   rv   r   r   )	directorymodel_idpartsmasked_directorys       r    get_model_masked_directoryr      sx     |u%% )  q))  a((w||E!Hh.>.>sE.J.JKKr"   	raw_inputc                    d}t          | t                    rZt                      }d| v r d |                     d          D             }n|                    |                                            nt          | t                    rt                      }| D ]u}t          |t                    r^d|v r3|                    d |                    d          D                        N|                    |                                           v|S )u   
    处理规则
    N,c                 6    g | ]}|                                 S r$   strip)r   ss     r    
<listcomp>z$convert_patterns.<locals>.<listcomp>5  s     >>>Aaggii>>>r"   c                 6    g | ]}|                                 S r$   r   )r   sss     r    r   z$convert_patterns.<locals>.<listcomp>=  s     "E"E"E"288::"E"E"Er"   )r   strlistsplitr   r   extend)r   outputr   s      r    convert_patternsr   -  s    F)S!! -)>>)=)=>>>FFMM)//++,,,,	It	$	$ - 	- 	-A!S!! -!88MM"E"E"E"E"EFFFFMM!'')),,,Mr"   r6   c                     |rd| }ndt          j        dd           }d|t          t                    d}| r|                    |            |S )u   
    填充header
    ztoken AISTUDIO_ACCESS_TOKENr6   )defaultzapplication/json)zContent-TypeAuthorizationzSDK-Version)ru   getenvr   r   rI   )paramstokenauthheaderss       r    header_fillr   C  su      IH	"92FFFHH*7|| G
  vNr"   c                 H   t          | dd          5 }|                                }ddd           n# 1 swxY w Y   t          j        d|t          j        t          j        z            }|r'|                    d                                          S t          d          )u   
    获取yaml
    rrb   )encodingNz^---\s*(.*?)\s*---r   u+   未找到两个 '---' 分隔的 YAML 内容)	rG   rH   r8   searchDOTALL	MULTILINEgroupr   
ValueError)rK   r}   rz   matchs       r    extract_yaml_blockr   U  s     
iw	/	/	/ 1&&((               I+Wbi",6NOOE H{{1~~##%%%FGGGs   488c                 L    t           j                            |           }|dk    S )u   
    判断文件名
    z	README.md)ru   rv   basename)rK   	file_names     r    is_readme_mdr   d  s$       ++I##r"   file_path_or_objreturnc                    t          | t          t          f          r(t          |           }|                                j        S t          | t
                    rt          |           S t          | t          j                  r_| 	                                }| 
                    dt          j                   | 	                                }| 
                    |           |S t          d          )z
    get size
    r   zCUnsupported type: must be string, Path, bytes, or io.BufferedIOBase)r   r   r   statst_sizebytesr   ioBufferedIOBasetellseekru   SEEK_END	TypeError)r   rK   current_positionry   s       r    get_file_sizer  l  s     "S$K00 
)**	~~''	$e	,	, 
#$$$	$b&7	8	8 	
+0022a---$$&&.///Q
 
 	
r"   r   [Calculating]Tbuffer_size_mbc           	      X   ddl m} t          |           }|dk    rSd}d}t          | t          t
          f          r-t          | t
                    r| nt          |           }|j        }d| d}|dz  dz  }t          j                    }	g }
 ||dd	d	d
||          }t          | t          t
          f          rt          | d          5 }|
                    |          x}r|
                    t          j        |                                                     |	                    |           |                    t          |                     |
                    |          x}ddd           n# 1 swxY w Y   |	                                }	|}njt          | t                    rc|	                    |            |	                                }	|
                    |	           t          |           }|                    |           nt          | t           j                  r| 
                    |          x}r|
                    t          j        |                                                     |	                    |           |                    t          |                     | 
                    |          x}|	                                }	|}n#|                                 t'          d          |                                 | |	||t          |
          |
dS )z
    calculate hash
    r   r   i   @Fz
Large Filez[Validating Hash for ]r   TB)r   r   r   dynamic_ncolsunitr   r   rB   Nz5Input must be str, Path, bytes or a io.BufferedIOBase)r   	file_hash	file_size
chunk_size
chunk_numschunk_hash_list)	tqdm.autor   r  r   r   r   rs   rE   rF   rG   rH   r   rJ   rI   r   r   r   r   r   r   )r   r  r   r   r   r  rs   rv   buffer_sizer
  r  progressr}   
byte_chunkfinal_chunk_sizes                  r    get_file_hashr    s    .//I%%%&d44 	'1 $(( (( D##-12B-C-C 9D3D333	 4'$.K  IOt  H "S$K00 E"D)) 	1Q !{ 3 33* 1&&w~j'A'A'K'K'M'MNNN  ,,,J000 !"{ 3 33* 1	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1 	1
 ''))	&	$e	,	, E)***''))	y)))/00())))	$b&7	8	8 E,11+>>>j 	-""7>*#=#=#G#G#I#IJJJZ(((OOC
OO,,, -11+>>>j 	- ''))	& 	CE E 	E NN -&/***  s   
BE55E9<E9)NN)Nr6   )r   r  T)=r)   rp   r   ru   r   r8   rd   rE   r   r   r   r   aistudio_sdkr   aistudio_sdk.errorsr   aistudio_sdk.configr   	functoolsr
   r  r   concurrent.futuresr   r   typingr   r   r   r   aistudio_sdk.constant.versionr   pathlibr   r   r   r,   r4   r>   r;   rO   rY   r`   rj   r   r   r   r   r   intboolr   r   r   r   r   r   r   r   r  r  r$   r"   r    <module>r     s}     



 				 				 				   2 2 2 2 2 2 2 2 2 2        2 2 2 2 2 2 3 3 3 3 3 3             ? ? ? ? ? ? ? ? 2 2 2 2 2 2 2 2 2 2 2 2 1 1 1 1 1 1          4   
 
 
   *    # # #
 
 
  *       0, , ,  "# # #& & &* (;).%)/ / /"&/"/ / / /d
 
 
c49n 5    ,   $H H H$ $ $
E#tUH*D$E 
# 
 
 
 
. %&.#'	I ICuh67ISMI }I 4.	I
 
I I I I I Ir"   