<s> """<STR_LIT>""" <EOL> import functools <EOL> from . compat import string_types , zip_longest <EOL> from . errors import DimensionalityError <EOL> from . util import to_units_container <EOL> def _replace_units ( original_units , values_by_name ) : <EOL> """<STR_LIT>""" <EOL> q = <NUM_LIT:1> <EOL> for arg_name , exponent in original_units . items ( ) : <EOL> q = q * values_by_name [ arg_name ] ** exponent <EOL> return to_units_container ( q ) <EOL> def _to_units_container ( a ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( a , string_types ) and '<STR_LIT:=>' in a : <EOL> return to_units_container ( a . split ( '<STR_LIT:=>' , <NUM_LIT:1> ) [ <NUM_LIT:1> ] ) , True <EOL> return to_units_container ( a ) , False <EOL> def _parse_wrap_args ( args ) : <EOL> defs_args = set ( ) <EOL> defs_args_ndx = set ( ) <EOL> dependent_args_ndx = set ( ) <EOL> unit_args_ndx = set ( ) <EOL> args_as_uc = [ _to_units_container ( arg ) for arg in args ] <EOL> for ndx , ( arg , is_ref ) in enumerate ( args_as_uc ) : <EOL> if arg is None : <EOL> continue <EOL> elif is_ref : <EOL> if len ( arg ) == <NUM_LIT:1> : <EOL> [ ( key , value ) ] = arg . items ( ) <EOL> if value == <NUM_LIT:1> and key not in defs_args : <EOL> defs_args . add ( key ) <EOL> defs_args_ndx . add ( ndx ) <EOL> args_as_uc [ ndx ] = ( key , True ) <EOL> else : <EOL> dependent_args_ndx . add ( ndx ) <EOL> else : <EOL> dependent_args_ndx . add ( ndx ) <EOL> else : <EOL> unit_args_ndx . add ( ndx ) <EOL> for ndx in dependent_args_ndx : <EOL> arg , is_ref = args_as_uc [ ndx ] <EOL> if not isinstance ( arg , dict ) : <EOL> continue <EOL> if not set ( arg . keys ( ) ) <= defs_args : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' % args [ ndx ] ) <EOL> def _converter ( ureg , values , strict ) : <EOL> new_values = list ( value for value in values ) <EOL> values_by_name = { } <EOL> for ndx in defs_args_ndx : <EOL> values_by_name [ args_as_uc [ ndx ] [ <NUM_LIT:0> ] ] = values [ ndx ] <EOL> new_values [ ndx ] = values [ ndx ] . _magnitude <EOL> for ndx in dependent_args_ndx : <EOL> new_values [ ndx ] = ureg . _convert ( values [ ndx ] . _magnitude , <EOL> values [ ndx ] . _units , <EOL> _replace_units ( args_as_uc [ ndx ] [ <NUM_LIT:0> ] , values_by_name ) ) <EOL> for ndx in unit_args_ndx : <EOL> if isinstance ( values [ ndx ] , ureg . Quantity ) : <EOL> new_values [ ndx ] = ureg . _convert ( values [ ndx ] . _magnitude , <EOL> values [ ndx ] . _units , <EOL> args_as_uc [ ndx ] [ <NUM_LIT:0> ] ) <EOL> else : <EOL> if strict : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' . format ( args_as_uc [ ndx ] [ <NUM_LIT:0> ] , new_values [ ndx ] ) ) <EOL> return new_values , values_by_name <EOL> return _converter <EOL> def wraps ( ureg , ret , args , strict = True ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( args , ( list , tuple ) ) : <EOL> args = ( args , ) <EOL> converter = _parse_wrap_args ( args ) <EOL> if isinstance ( ret , ( list , tuple ) ) : <EOL> container , ret = True , ret . __class__ ( [ _to_units_container ( arg ) for arg in ret ] ) <EOL> else : <EOL> container , ret = False , _to_units_container ( ret ) <EOL> def decorator ( func ) : <EOL> assigned = tuple ( attr for attr in functools . WRAPPER_ASSIGNMENTS if hasattr ( func , attr ) ) <EOL> updated = tuple ( attr for attr in functools . WRAPPER_UPDATES if hasattr ( func , attr ) ) <EOL> @ functools . wraps ( func , assigned = assigned , updated = updated ) <EOL> def wrapper ( * values , ** kw ) : <EOL> new_values , values_by_name = converter ( ureg , values , strict ) <EOL> result = func ( * new_values , ** kw ) <EOL> if container : <EOL> out_units = ( _replace_units ( r , values_by_name ) if is_ref else r <EOL> for ( r , is_ref ) in ret ) <EOL> return ret . __class__ ( res if unit is None else ureg . Quantity ( res , unit ) <EOL> for unit , res in zip ( out_units , result ) ) <EOL> if ret [ <NUM_LIT:0> ] is None : <EOL> return result <EOL> return ureg . Quantity ( result , <EOL> _replace_units ( ret [ <NUM_LIT:0> ] , values_by_name ) if ret [ <NUM_LIT:1> ] else ret [ <NUM_LIT:0> ] ) <EOL> return wrapper <EOL> return decorator <EOL> def check ( ureg , * args ) : <EOL> """<STR_LIT>""" <EOL> dimensions = [ ureg . get_dimensionality ( dim ) for dim in args ] <EOL> def decorator ( func ) : <EOL> assigned = tuple ( attr for attr in functools . WRAPPER_ASSIGNMENTS if hasattr ( func , attr ) ) <EOL> updated = tuple ( attr for attr in functools . WRAPPER_UPDATES if hasattr ( func , attr ) ) <EOL> @ functools . wraps ( func , assigned = assigned , updated = updated ) <EOL> def wrapper ( * values , ** kwargs ) : <EOL> for dim , value in zip_longest ( dimensions , values ) : <EOL> if dim and value . dimensionality != dim : <EOL> raise DimensionalityError ( value , '<STR_LIT>' , <EOL> value . dimensionality , dim ) <EOL> return func ( * values , ** kwargs ) <EOL> return wrapper <EOL> return decorator </s>
