B
    T\                 @   s(  d dl mZmZm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 ddlmZ ddlmZ y
e
jZW nX ek
r   ye
d gjdd	 W n* ek
r Z zeeZW ddZ[X Y nX Y nX d
d ZdddZdd Zdd Zdd ZdddZd ddZej fddZ!dd Z"dS )!    )absolute_importdivisionprint_functionN)frequenciesconcat   )Array   )HighLevelGraph   )axisc             C   s    dt t| kr|  S | S d S )NZcupy)strtypeget)x r   /lib/python3.7/site-packages/dask/array/utils.pynormalize_to_array   s    r   Fc             K   sp   t | } t |}t| dd dkr8tj| |fd|i|S |rd| j|jkobtdd t| j|jD S | |k S )NdtypeO	equal_nanc             s   s.   | ]&\}}t |rt |n||kV  qd S )N)npisnan).0abr   r   r   	<genexpr>&   s   zallclose.<locals>.<genexpr>)r   getattrr   allcloseshapeallzipZflat)r   r   r   kwargsr   r   r   r      s    r   c             C   s$   dd }t | j|dt |j|dkS )Nc             S   s   t | tr| dddfS | S d S )N)
isinstancer   )kr   r   r   key,   s    
zsame_keys.<locals>.key)r&   )sorteddask)r   r   r&   r   r   r   	same_keys+   s    r)   c             C   s   | j od| j kS )Nr   )r   )r   r   r   r   
_not_empty4   s    r*   c             C   sZ   t | tsdS tdd | jD s&ttt| j }dd |	 D }|rVt|dS )z4 Check that graph is well named and non-overlapping Nc             s   s   | ]}t |ttfV  qd S )N)r$   tupler   )r   r%   r   r   r   r   =   s    z_check_dsk.<locals>.<genexpr>c             S   s   i | ]\}}|d kr||qS )r   r   )r   r%   vr   r   r   
<dictcomp>?   s    z_check_dsk.<locals>.<dictcomp>)
r$   r
   r    ZlayersAssertionErrorr   r   Zdictsvaluesitems)ZdskZfreqsZnon_oner   r   r   
_check_dsk8   s    
r1   Tc             C   sZ   xTt | |D ]F\}}t|s(t|rF|rRt|t|ksRtq||kstqW d S )N)r!   mathr   r.   )r   r   	check_nanZaaZbbr   r   r   assert_eq_shapeC   s
    r4   c             K   s@  | }|}t | tr| jd k	s t| j}|r4t| j | jdd} t| drR|  } t| dsjt	j
| dd} t| r| j|jkst|rt|j| jdd n$t| dst	j
| dd} t| dd }t |tr\|jd k	st|j}|rt|j |jdd}t|dst	j
|dd}t|dr&| }t|rB|j|jksBt|rt|j|jdd n&t|dsvt	j
|dd}t|dd }t|t|krtt| t| }	td	tj tj|	 y*| j|jkstt| |f|std
S  tk
r
   Y nX | |k}
t |
t	jr2|
 s<tn
|
s<td
S )Nsync)Z	schedulertodenser   r   )r   F)r3   zstring repr are differentT)r$   r   r   r.   r1   r(   Zcomputehasattrr6   r   arrayr*   r4   r   r   r   difflibZndiff
splitlinesoslinesepjoinr   	TypeErrorZndarrayr    )r   r   Zcheck_shapeZcheck_graphr"   Z
a_originalZ
b_originalZadtZbdtZdiffcr   r   r   	assert_eqL   sf    








r@   c                s0   t  fdd|D r$tj |dS dd S dS )zkLike functools.wraps, but safe to use even if wrapped is not a function.

    Only needed on Python 2.
    c             3   s   | ]}t  |V  qd S )N)r7   )r   attr)wrappedr   r   r      s    zsafe_wraps.<locals>.<genexpr>)assignedc             S   s   | S )Nr   )r   r   r   r   <lambda>   s    zsafe_wraps.<locals>.<lambda>N)r    	functoolswraps)rB   rC   r   )rB   r   
safe_wraps   s    rG   c                sr   t | ttfr$t fdd| D S t | tjs<td|  |   k sN|  kr^td|  f | dk rn|  7 } | S )z% Validate an input to axis= keywords c             3   s   | ]}t | V  qd S )N)validate_axis)r   Zax)ndimr   r   r      s    z validate_axis.<locals>.<genexpr>z%Axis value must be an integer, got %sz2Axis %d is out of bounds for array of dimension %dr   )r$   r+   listnumbersZIntegralr>   	AxisError)r   rI   r   )rI   r   rH      s    rH   )F)T)TT)#Z
__future__r   r   r   r9   rE   r2   rK   r;   Znumpyr   Ztoolzr   r   Zcorer   Zhighlevelgraphr
   rL   AttributeErrorr8   sum	Exceptioner   r   r   r)   r*   r1   r4   r@   WRAPPER_ASSIGNMENTSrG   rH   r   r   r   r   <module>   s2   
 
	
	
?