1# ----------------------------------------------------------- 2# Copyright (C) 2009 StatPro Italia s.r.l. 3# 4# StatPro Italia 5# Via G. B. Vico 4 6# I-20123 Milano 7# ITALY 8# 9# phone: +39 02 96875 1 10# fax: +39 02 96875 605 11# 12# email: info@riskmap.net 13# 14# This program is distributed in the hope that it will be 15# useful, but WITHOUT ANY WARRANTY; without even the 16# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR 17# PURPOSE. See the license for more details. 18# ----------------------------------------------------------- 19# 20# Author: Enrico Sirola <enrico.sirola@statpro.com> 21 22"""internal helpers""" 23 24import ctypes as _ct 25from drmaa.wrappers import * 26from drmaa.errors import error_buffer 27import drmaa.const as const 28 29_BUFLEN=const.ATTR_BUFFER 30 31try: 32 import namedtuple as _nt 33except ImportError: # pre 2.6 behaviour 34 import nt as _nt 35 36class BoolConverter(object): 37 """Helper class to convert to/from bool attributes.""" 38 def __init__(self, true='y', false='n'): 39 self.true=true 40 self.false=false 41 def to_drmaa(self, value): 42 if value: return self.true 43 else: return self.false 44 def from_drmaa(self, value): 45 if value == self.true: 46 return True 47 else: 48 return False 49 50class IntConverter(object): 51 """Helper class to convert to/from float attributes.""" 52 @staticmethod 53 def to_drmaa(value): 54 return str(value) 55 56 @staticmethod 57 def from_drmaa(value): 58 return int(value) 59 60class SessionStringAttribute(object): 61 62 def __init__(self, drmaa_function): 63 self._f = drmaa_function 64 def __get__(self, *args): 65 buf = _ct.create_string_buffer(_BUFLEN) 66 c(self._f, buf, _ct.sizeof(buf)) 67 return buf.value 68 69Version = _nt.namedtuple("Version", "major minor") 70Version.__str__ = lambda x: "%s.%s" % (x.major, x.minor) 71#Version.__doc__ = """\ 72#An object representing the DRMAA version. 73# 74#major and minor attributes are int. For DRMAA 1.0, major == 1 and minor == 0. 75#""" 76 77class SessionVersionAttribute(object): 78 """A Version attribute.""" 79 def __get__(self, *args): 80 major = _ct.c_uint(10) 81 minor = _ct.c_uint(10) 82 c(drmaa_version, _ct.byref(major), _ct.byref(minor)) 83 return Version(major.value, minor.value) 84 85class Attribute(object): 86 """A DRMAA attribute, to be managed with scalar C DRMAA attribute management functions.""" 87 def __init__(self, name, type_converter=None): 88 """\ 89Attribute constructor. 90 91:Parameters: 92 `name` : string 93 name of the attribute to be managed, as seen by the underlying C DRMAA 94 `type_converter` 95 a converter to translate attribute values to/from the underlying 96 implementation. See BoolConverter for an example. 97""" 98 self.name = name 99 self.converter = type_converter 100 def __set__(self, instance, value): 101 if self.converter: 102 v = self.converter.to_drmaa(value) 103 else: 104 v = value 105 c(drmaa_set_attribute, instance, self.name, v) 106 def __get__(self, instance, _): 107 attr_buffer = create_string_buffer(const.ATTR_BUFFER) 108 c(drmaa_get_attribute, instance, self.name, 109 attr_buffer, sizeof(attr_buffer)) 110 if self.converter: 111 return self.converter.from_drmaa(attr_buffer.value) 112 else: 113 return attr_buffer.value 114 115class VectorAttribute(object): 116 """\ 117A DRMAA attribute representing a list. 118 119To be managed with vector C DRMAA attribute management functions.""" 120 def __init__(self, name): 121 self.name = name 122 def __set__(self, instance, value): 123 c(drmaa_set_vector_attribute, instance, self.name, string_vector(value)) 124 def __get__(self, instance, _): 125 return list(vector_attribute_iterator( 126 instance, self.name)) 127 128class DictAttribute(object): 129 """\ 130A DRMAA attribute representing a python dict. 131 132To be managed with vector C DRMAA attribute management functions.""" 133 def __init__(self, name): 134 self.name = name 135 def __set__(self, instance, value): 136 v = [ "%s=%s" % (k, v) for (k, v) in value.iteritems() ] 137 c(drmaa_set_vector_attribute, instance, self.name, string_vector(v)) 138 def __get__(self, instance, _): 139 x = [ i.split('=', 1) for i in list(vector_attribute_iterator( 140 instance, self.name)) ] 141 return dict(x) 142 143 144def attributes_iterator(attributes): 145 try: 146 buf = create_string_buffer(const.ATTR_BUFFER) 147 while drmaa_get_next_attr_value(attributes, buf, sizeof(buf))\ 148 != const.NO_MORE_ELEMENTS: 149 yield buf.value 150 except: 151 drmaa_release_attr_values(attributes) 152 raise 153 else: 154 drmaa_release_attr_values(attributes) 155 156def adapt_rusage(rusage): 157 "transform a rusage data structure into a dict" 158 rv = dict() 159 for attr in attributes_iterator(rusage.contents): 160 k, v = attr.split('=') 161 rv[k] = v 162 return rv 163 164def vector_attribute_iterator(jt, attr_name): 165 avalues = pointer(POINTER(drmaa_attr_values_t)()) 166 c(drmaa_get_vector_attribute, jt, attr_name, avalues) 167 return attributes_iterator(avalues.contents) 168 169def attribute_names_iterator(): 170 attrn_p = pointer(POINTER(drmaa_attr_names_t)()) 171 c(drmaa_get_attribute_names, attrn_p) 172 try: 173 name = create_string_buffer(_BUFLEN) 174 while drmaa_get_next_attr_name(attrn_p.contents, name, _BUFLEN)\ 175 != const.NO_MORE_ELEMENTS: 176 yield name.value 177 except: 178 drmaa_release_attr_names(attrn_p.contents) 179 raise 180 else: 181 drmaa_release_attr_names(attrn_p.contents) 182 183def vector_attribute_names_iterator(): 184 attrn_p = pointer(POINTER(drmaa_attr_names_t)()) 185 c(drmaa_get_vector_attribute_names, attrn_p) 186 try: 187 name = create_string_buffer(_BUFLEN) 188 while drmaa_get_next_attr_name(attrn_p.contents, name, _BUFLEN)\ 189 != const.NO_MORE_ELEMENTS: 190 yield name.value 191 except: 192 drmaa_release_attr_names(attrn_p.contents) 193 raise 194 else: 195 drmaa_release_attr_names(attrn_p.contents) 196 197def run_bulk_job(jt, start, end, incr=1): 198 jids = pointer(POINTER(drmaa_job_ids_t)()) 199 try: 200 c(drmaa_run_bulk_jobs, jids, jt, start, end, incr) 201 jid = create_string_buffer(_BUFLEN) 202 while drmaa_get_next_job_id(jids.contents, jid, _BUFLEN)\ 203 != const.NO_MORE_ELEMENTS: 204 yield jid.value 205 except: 206 drmaa_release_job_ids(jids.contents) 207 raise 208 else: 209 drmaa_release_job_ids(jids.contents) 210 211def c(f, *args): 212 """An helper function wrapping calls to the C DRMAA functions with error managing code.""" 213 return f(*(args + (error_buffer, sizeof(error_buffer)))) 214 215def string_vector(v): 216 vlen = len(v) 217 values = (STRING * (vlen + 1))() 218 for i, el in enumerate(v): 219 values[i] = STRING(el) 220 values[vlen] = STRING() 221 return values 222 223def attribute_setter(obj, attribute_name): 224 "returns a drmaa attribute setter" 225 def f(value): 226 "setter for %s" % attribute_name 227 c(drmaa_set_attribute, obj, attribute_name, value) 228 f.__name__ = 'set_'+attribute_name 229 return f 230 231def attribute_getter(obj, attribute_name): 232 "returns a drmaa attribute setter" 233 def f(): 234 "getter for %s" % attribute_name 235 attr_buffer = create_string_buffer(1024) 236 c(drmaa_get_attribute, obj, attribute_name, attr_buffer, sizeof(attr_buffer)) 237 return attr_buffer.value 238 f.__name__ = 'get_'+attribute_name 239 return f 240