rosh.impl.bagy

1 # Software License Agreement (BSD License) 2 # 3 # Copyright (c) 2010, Willow Garage, Inc. 4 # All rights reserved. 5 # 6 # Redistribution and use in source and binary forms, with or without 7 # modification, are permitted provided that the following conditions 8 # are met: 9 # 10 # * Redistributions of source code must retain the above copyright 11 # notice, this list of conditions and the following disclaimer. 12 # * Redistributions in binary form must reproduce the above 13 # copyright notice, this list of conditions and the following 14 # disclaimer in the documentation and/or other materials provided 15 # with the distribution. 16 # * Neither the name of Willow Garage, Inc. nor the names of its 17 # contributors may be used to endorse or promote products derived 18 # from this software without specific prior written permission. 19 # 20 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 21 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 22 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS 23 # FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE 24 # COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, 25 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, 26 # BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 27 # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 28 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 29 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN 30 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 31 # POSSIBILITY OF SUCH DAMAGE. 32 # 33 # Revision $Id: bagy.py 11681 2010-10-22 07:08:51Z kwc $ 34 35 from __future__ import with_statement 36 37 """ 38 YAML-based bag formats. 39 """ 40 41 #TODO: convert to ABCs 42 #TODO: orderedkey bagy ? 43 #TODO: multimap bagy? 44 45 import collections 46 import yaml 47 48 import roslib.message 49 50 import rosh 51 from rosh.impl.exceptions import ROSHException 52 from rosh.impl.topic import TopicNS 53

54 -def create(msg_repr, type_, filename=None):

55 """ 56 Create single message from YAML/dictionary representation. 57 58 @param type_: Message class 59 @type type_: class (Message subclass) 60 @param msg_repr: Message in string or dictionary representation 61 @type msg_repr: str or dict 62 @param filename: Name of file associated with msg_repr for debugging. 63 @type filename: str 64 """ 65 if type(msg_repr) == str: 66 msg_repr = yaml.load(msg_repr) 67 if type(msg_repr) != dict: 68 if filename: 69 raise ValueError("msg_repr file [%s] does not contain a YAML dictionary"%filename) 70 else: 71 raise ValueError("msg_repr must contain a dictionary") 72 m = type_() 73 # TODO: add new API to roslib.message that doesn't assume a rostopic-like API source 74 roslib.message.fill_message_args(m, [msg_repr]) 75 return m

76 77 ZERO_STATE = 0 78 RECORD_STATE = 1 79 STOP_STATE = 2 80

81 -class _Bagy(object):

82

83 - def __init__(self, name, stream, type_):

84 """ 85 ctor. 86 """ 87 self.name = name 88 self.type_ = type_ 89 self._stream = stream 90 self._buff = None 91 self._yaml_iter = None 92 self.closed = False 93 94 # for recording 95 self._state = ZERO_STATE 96 self._record = None

97

98 - def __enter__(self):

99 return self

100

101 - def __exit__(self, *args):

102 self.close()

103

104 - def __iter__(self):

105 for doc in yaml.load_all(self._stream): 106 if doc is not None: 107 yield create(doc, self.type_)

108

109 - def write(self, msg):

110 """ 111 Write message to bagy 112 @param msg: messsage to add 113 @type msg: self.type_ 114 """ 115 if not isinstance(msg, roslib.message.Message): 116 raise ValueError("msg is not a message instance") 117 self._stream.write(roslib.message.strify_message(msg) + '\n---\n')

118

119 - def read(self):

120 """ 121 Read *all* messages from bagy. 122 123 @return: messages 124 @rtype: [self.type_] 125 """ 126 return [create(d, self.type_) for d in yaml.load_all(self._stream) if d is not None]

127

128 - def next(self):

129 """ 130 Get the next message in the bagy 131 132 @return: Message instance 133 @rtype: self.type_ 134 """ 135 if self._yaml_iter is None: 136 self._yaml_iter = yaml.load_all(self._stream) 137 v = self._yaml_iter.next() 138 if v is None: 139 return None 140 return create(v, self.type_)

141

142 - def close(self):

143 """ 144 Close resources associated with Bagy 145 """ 146 self.closed = True 147 if self._state == RECORD_STATE: 148 self.stop() 149 self._state = ZERO_STATE 150 self._record = None 151 152 self._stream.close() 153 if self._yaml_iter is not None: 154 self._yaml_iter.close()

155

156 - def record(self, topic):

157 if self.closed: 158 raise ROSHException("closed") 159 if self._record is not None: 160 raise ROSHException("already recording topic %s"%(self._record._name)) 161 162 if isinstance(topic, TopicNS): 163 pass 164 elif type(topic) == str: 165 ctx = rosh.get_default_plugin_context().ctx 166 topic = ctx.topics[topic] 167 168 # TODO: clean this up once I add proper subscribe() to topic 169 topic._add_subscriber_callback(self.write) 170 self._record = topic 171 self._state = RECORD_STATE

172

173 - def start(self):

174 if self.closed: 175 raise ROSHException("closed") 176 if self._record is None: 177 raise ROSHException("not recording") 178 if self._state == STOP_STATE: 179 self._record._add_subscriber_callback(self.write) 180 self._state = RECORD_STATE

181

182 - def stop(self):

183 if self._record is None: 184 raise ROSHException("not recording") 185 186 if self._state == RECORD_STATE: 187 self._record._remove_subscriber_callback(self.write) 188 self._state = STOP_STATE

189 190

191 -class _MapBagy(collections.Mapping):

192 """ 193 Bagy where messages are indexed by labels (keys). 194 """

195 - def __init__(self, name, stream, type_):

196 """ 197 ctor. 198 """ 199 self.name = name 200 self.type_ = type_ 201 self._stream = stream 202 self._buff = None 203 self._db = None

204

205 - def __enter__(self):

206 return self

207

208 - def __exit__(self, *args):

209 self.close()

210

211 - def __iter__(self):

212 if self._db is None: 213 self._load_db() 214 for key, d in self._db.iteritems(): 215 yield key, create(d, self.type_)

216 217 # collections.Mapping

218 - def __contains__(self, key):

219 if self._db is None: 220 self._load_db() 221 return key in self._db

222 223 # collections.Mapping

224 - def __eq__(self, other):

225 # easy but expensive way 226 return self.items().__eq__(other.items())

227 228 # collections.Mapping

229 - def __ne__(self, other):

230 # easy but expensive way 231 return self.items().__ne__(other.items())

232 233 # collections.Mapping

234 - def keys(self):

235 if self._db is None: 236 self._load_db() 237 return self._db.keys()

238 239 # collections.Mapping

240 - def items(self):

241 return list(self.iteritems())

242 243 # collections.Mapping

244 - def values(self):

245 return list(self.itervalues())

246

247 - def iterkeys(self):

248 if self._db is None: 249 self._load_db() 250 return self._db.iterkeys()

251

252 - def iteritems(self):

253 return self.__iter__()

254

255 - def itervalues(self):

256 if self._db is None: 257 self._load_db() 258 for d in self._db.itervalues(): 259 yield create(d, self.type_)

260

261 - def write(self, *args, **kwds):

262 """ 263 write message into bagy. Must either pass in label and msg instance:: 264 265 b.write('key', m) 266 267 or, pass in Python keywords, where keys are the labels:: 268 269 b.write(key1=m1, key2=m2) 270 """ 271 if args: 272 key, msg = args 273 if not isinstance(key, basestring): 274 raise TypeError("label must be a string") 275 if not isinstance(msg, roslib.message.Message): 276 raise ValueError("msg is not a message instance") 277 self._stream.write('%s:%s\n'%(key, 278 roslib.message.strify_message(msg, indent=' '))) 279 elif kwds: 280 for key, msg in kwds.iteritems(): 281 if not isinstance(key, basestring): 282 raise TypeError("label must be a string") 283 if not isinstance(msg, roslib.message.Message): 284 raise ValueError("msg is not a message instance") 285 self._stream.write('%s:%s\n'%(key, 286 roslib.message.strify_message(msg, indent=' '))) 287 else: 288 raise TypeError("must provide an arguments or keywords to write()")

289

290 - def _load_db(self):

291 """ 292 Load YAML dictionary into bagy instance. We don't convert the 293 dicts to messages until they are requested. 294 """ 295 db = {} 296 for d in yaml.load_all(self._stream): 297 db.update(d) 298 self._stream.close() 299 self._db = db

300

301 - def __setitem__(self, key, value):

302 """ 303 Add item to label bagy as a map, e.g.:: 304 305 b = keybagy('foo.bagy', 'w') 306 b[key] = msg 307 """ 308 self.write(key, value)

309 310 # collections.Mapping

311 - def get(self, key, default=None):

312 """ 313 Get value from bagy, with optional default, e.g.:: 314 315 b = keybagy('foo.bagy', 'r', msg.std_msgs.String) 316 msg = b.(key, altmsg) 317 """ 318 if self._db is None: 319 self._load_db() 320 if key in self._db: 321 return create(self._db[key], self.type_) 322 else: 323 return default

324 325 # collections.Mapping

326 - def __len__(self, key):

327 if self._db is None: 328 self._load_db() 329 return self._db.__len__()

330 331 # collections.Mapping

332 - def __getitem__(self, key):

333 """ 334 Access label bagy as a map, e.g.:: 335 336 b = keybagy('foo.bagy', 'r', msg.std_msgs.String) 337 msg = b[key] 338 """ 339 if self._db is None: 340 self._load_db() 341 return create(self._db[key], self.type_)

342

343 - def close(self):

344 self._stream.close() 345 if self._db is not None: 346 del self._db

347

348 - def read(self):

349 """ 350 Read *all* messages from label bagy. This exhausts the bagy 351 and future methods that access data from the bagy will fail. 352 353 @return: messages 354 @rtype: dict{str: self.type_} 355 """ 356 if self._db is None: 357 self._load_db() 358 359 retval = {} 360 for k, v in self._db.iteritems(): 361 retval[k] = create(v, self.type_) 362 del self._db 363 self._db = None 364 return retval

365 366 367

368 -class Bagy(_Bagy):

369 - def __init__(self, filename, mode='r', type_=None):

370 """ 371 @param filename: path to file 372 @type filename: str 373 @param mode: open mode. Valid modes are 'r' (read), 'w' (write), and 'a' (append). 374 @type mode: str 375 @param type_: Message type stored in bagy 376 @type type_: type(roslib.message.Message) 377 """ 378 if mode in ['r', 'w', 'a']: 379 if mode == 'r' and type_ is None: 380 raise ValueError("type_ must be specified when reading from bagy files") 381 self.mode = mode 382 super(Bagy, self).__init__(filename, open(filename, mode), type_) 383 else: 384 raise ValueError("mode is invalid")

385

386 -class MapBagy(_MapBagy):

387 - def __init__(self, filename, mode, type_=None):

388 """ 389 @param filename: path to file 390 @type filename: str 391 @param mode: open mode. Valid modes are 'r' (read), 'w' (write), and 'a' (append). 392 @type mode: str 393 @param type_: Message type stored in key bagy 394 @type type_: type(roslib.message.Message) 395 """ 396 if mode in ['r', 'w', 'a']: 397 if mode == 'r' and type_ is None: 398 raise ValueError("type_ must be specified when reading from bagy files") 399 self.mode = mode 400 super(MapBagy, self).__init__(filename, open(filename, mode), type_) 401 else: 402 raise ValueError("mode is invalid")

403

Source Code for Module rosh.impl.bagy