3
���h3+ � @ s� d dl Z d dlmZ d dlZddlmZ ddlmZm Z ddl
m Z ddl m
Z
mZmZmZ ddlmZmZmZmZ d d
lmZmZ d d lmZmZm Z d d
dd d�Z G dd� de�Z dd� Z dd� Z dd� Z!dd� Z"d!d d �Z#G d d � d e�Z$dS )"� N)�partial� )�Variable)�pop_to�
cf_encoder)�indexing)�FrozenOrderedDict� NDArrayMixin�close_on_error�
is_remote_uri)� iteritems�
basestring� OrderedDict�PY3� )�WritableCFDataStore�robust_getitem)�encode_nc3_attr_value�encode_nc3_variable�maybe_convert_to_char_array�native�big�little)�=�>�<�|c @ s* e Zd Zd dd�Zedd� �Zdd� ZdS )
�NetCDF4ArrayWrapperFc C s || _ || _d S )N)�array� is_remote)�selfr r � r! �8/tmp/pip-build-5_djhm0z/xray/xarray/backends/netCDF4_.py�__init__ s z NetCDF4ArrayWrapper.__init__c C s | j j}|tkrtjd�}|S )N�O)r �dtype�str�np)r r% r! r! r" r% s
zNetCDF4ArrayWrapper.dtypec C s� | j rtttd�}ntj}y|| j|�}W n< tk
rf d}tsZdd l }|d|j
� 7 }t|��Y nX | j dkr|t j
|�}|S )N)Zcatchz�The indexing operation you are attempting to perform is not valid on netCDF4.Variable object. Try loading your data into memory first by calling .load().r z
Original traceback:
)r r r � RuntimeError�operator�getitemr �
IndexErrorr � traceback�
format_exc�ndimr' Zasscalar)r �keyr* �data�msgr, r! r! r" � __getitem__( s
z NetCDF4ArrayWrapper.__getitem__N)F)�__name__�
__module__� __qualname__r# �propertyr% r2 r! r! r! r" r s
r c C s� | j jdkr(t| �dkr"| jd�} t}nT| j jdkr`t| j| j�\}}t||| j | j
�} | j }n | j jd
krt| j }nt d ��| |fS ) N�Ur r$ �S�i�u�f�cz'cannot infer dtype for netCDF4 variable)r9 r: r; r< ) r% �kind�len�astyper& r r0 �dimsr �attrs�encoding�
ValueError)�varr% r0 r@ r! r! r" �_nc4_values_and_dtypeD s
rE c C s� |t d ddg�kr| S t|t�s(td��|jd�jd�}x`|D ]X}y| j| } W q> tk
r� } z(|dkrv| j|�} nt d| |��W Y d d }~X q>X q>W | S d S )N� �/z group must be a string or None�rzgroup not found: %s)
�set�
isinstancer
rC �strip�split�groups�KeyErrorZ createGroup�IOError)�ds�group�mode�pathr/ �er! r! r" �
_nc4_groupW s
$rU c C s* | j jdkr&d|kr&tj|d �|d< d S )Nr8 �
_FillValue)r% r= r' Zstring_)r0 �
attributesr! r! r" �_ensure_fill_value_validm s rX c C s` | j jdkrB| jj| j jd��}t| j|| j| j�} | jj dd � | jj
dd�dk r\t d��| S )Nr r �endianr z]Attempt to write non-native endian type, this is not supported by the netCDF4 python library.)r r ) r% � byteorderr0 r? Z newbyteorderr r@ rA rB �pop�get�NotImplementedError)rD r0 r! r! r" �_force_native_endiannesst s r^ FT�netCDF4c s� | j j� }tddg�}tdddddg�� |r6� jd� |jd�d k rd|jd| j�| jkrd| rd|d= x|D ]}||krj||= qjW |r�� fd d
�|D �}|r�td ||f ��n x t|�D ]}|� kr�||= q�W |S ) N�source�original_shape�zlib� complevel�
fletcher32�
contiguous�
chunksizes�least_significant_digitc s g | ]}|� kr|�qS r! r! )�.0�k)�valid_encodingsr! r" �
<listcomp>� s z)_extract_nc4_encoding.<locals>.<listcomp>z2unexpected encoding parameters for %r backend: %r)rB �copyrI �addr\ �shaperC �list)�variable�raise_on_invalidZlsd_okay�backendrB Z safe_to_dropri �invalidr! )rj r" �_extract_nc4_encoding� s.
rt c sp e Zd ZdZd� fdd� Zd d
� Zd d � Zd
d� Zdd� Zdd� Z dd� Z
d dd�Z � fdd�Z dd� Z
� ZS ) �NetCDF4DataStorez�Store for reading and writing data via the Python-NetCDF4 library.
This store supports NetCDF3, NetCDF4 and OpenDAP datasets.
rH �NETCDF4NTFc sv dd l } |d krd}| j||||||d�}
t|
�� t|
||�| _W d Q R X || _t|�| _|| _t t
| �j |� d S )Nr rv )rR �clobber�diskless�persist�format) r_ ZDatasetr
rU rP rz r r � _filename�superru r# ) r �filenamerR rz rQ �writerrw rx ry Znc4rP )� __class__r! r" r# � s
zNetCDF4DataStore.__init__c s� � j d� � j}tjt� | j��}t� fdd�� j� D ��}t||� i }� j � }|d k rd|j
|� � j � }|d k r�|dkr�d|d< d |d< nd|d< t |�|d< t
||d� | j|d< � j|d < t||||�S )
NFc 3 s | ]}|� j |�fV qd S )N)� getncattr)rh ri )rD r! r" � <genexpr>� s z7NetCDF4DataStore.open_store_variable.<locals>.<genexpr>re Trf rg r` ra )�set_auto_maskandscale�
dimensionsr ZLazilyIndexedArrayr r r �ncattrsrX �filters�update�chunking�tupler r{ rn r )r rD r� r0 rW rB r� r� r! )rD r" �open_store_variable� s,
z$NetCDF4DataStore.open_store_variablec s t � fdd�t� jj�D ��S )Nc 3 s | ]\}}|� j |�fV qd S )N)r� )rh ri �v)r r! r" r� � s z1NetCDF4DataStore.get_variables.<locals>.<genexpr>)r r rP � variables)r r! )r r" �
get_variables� s z NetCDF4DataStore.get_variablesc s t � fdd�� jj� D ��S )Nc 3 s | ]}|� j j|�fV qd S )N)rP r� )rh ri )r r! r" r� � s z-NetCDF4DataStore.get_attrs.<locals>.<genexpr>)r rP r� )r r! )r r" � get_attrs� s zNetCDF4DataStore.get_attrsc C s t dd� t| jj�D ��S )Nc s s | ]\}}|t |�fV qd S )N)r> )rh ri r� r! r! r" r� � s z2NetCDF4DataStore.get_dimensions.<locals>.<genexpr>)r r rP r� )r r! r! r" �get_dimensions� s z NetCDF4DataStore.get_dimensionsc C s | j j||d� d S )N)�size)rP ZcreateDimension)r �name�lengthr! r! r" �
set_dimension� s z NetCDF4DataStore.set_dimensionc C s$ | j dkrt|�}| jj||� d S )Nrv )rz r rP � setncattr)r r/ �valuer! r! r" �
set_attribute� s
z NetCDF4DataStore.set_attributec C s� |j j� }t|�}| jdkr*t|�\}}nt|�}|j}| j|� |jdd �}|dkrZd }t ||d�}| j
j |||j |j
dd�|j
dd �|j
d
d �|j
d d�|j
d
d�|j
d�d|j
d�|d� }|jd� x t|�D ]\} }
|j| |
� q�W ||jfS )Nrv rV rF � )rq rb Frc � �shuffleTrd re rf r rg ) Zvarname�datatyper� rb rc r� rd re rf rY rg �
fill_value)rF r� )rA rl r^ rz rE r r% Zset_necessary_dimensionsr[ rt rP ZcreateVariabler@ r\ r� r r� r0 ) r r� rp Zcheck_encodingrA r� r� rB Znc4_varri r� r! r! r" �prepare_variable� s:
z!NetCDF4DataStore.prepare_variablec s t t| �j� | jj� d S )N)r| ru �syncrP )r )r r! r" r� s zNetCDF4DataStore.syncc C s( | j }x|jd k r|j}qW |j� d S )N)rP �parent�close)r rP r! r! r" r� s
zNetCDF4DataStore.close)rH rv NNTFF)F)r3 r4 r5 �__doc__r# r� r� r� r� r� r� r� r� r� �
__classcell__r! r! )r r" ru � s
* ru )FTr_ )%r) � functoolsr Znumpyr' rF r Z conventionsr r �corer Z
core.utilsr r r
r Z
core.pycompatr r
r r �commonr r Znetcdf3r r r Z_endian_lookupr rE rU rX r^ rt ru r! r! r! r" �<module> s* +
|