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