source: trunk/CifFile/StarFile.py @ 469

Last change on this file since 469 was 469, checked in by toby, 11 years ago

rework phase import
  • Property svn:executable set to *
File size: 44.6 KB
Line 
1"""
21.This Software copyright \u00A9 Australian Synchrotron Research Program Inc, ("ASRP").
3
42.Subject to ensuring that this copyright notice and licence terms
5appear on all copies and all modified versions, of PyCIFRW computer
6code ("this Software"), a royalty-free non-exclusive licence is hereby
7given (i) to use, copy and modify this Software including the use of
8reasonable portions of it in other software and (ii) to publish,
9bundle and otherwise re-distribute this Software or modified versions
10of this Software to third parties, provided that this copyright notice
11and terms are clearly shown as applying to all parts of software
12derived from this Software on each occasion it is published, bundled
13or re-distributed.  You are encouraged to communicate useful
14modifications to ASRP for inclusion for future versions.
15
163.No part of this Software may be sold as a standalone package.
17
184.If any part of this Software is bundled with Software that is sold,
19a free copy of the relevant version of this Software must be made
20available through the same distribution channel (be that web server,
21tape, CD or otherwise).
22
235.It is a term of exercise of any of the above royalty free licence
24rights that ASRP gives no warranty, undertaking or representation
25whatsoever whether express or implied by statute, common law, custom
26or otherwise, in respect of this Software or any part of it.  Without
27limiting the generality of the preceding sentence, ASRP will not be
28liable for any injury, loss or damage (including consequential loss or
29damage) or other loss, loss of profits, costs, charges or expenses
30however caused which may be suffered, incurred or arise directly or
31indirectly in respect of this Software.
32
336. This Software is not licenced for use in medical applications.
34"""
35
36from types import *
37from urllib import *         # for arbitrary opening
38import re
39import copy
40class StarList(list):
41    pass
42
43# Because DDLm makes a tuple from a tuple...
44class StarTuple(tuple):
45    def __new__(cls,*arglist):
46        return tuple.__new__(cls,arglist)
47
48class StarDict(dict):
49    pass
50
51class LoopBlock:
52    def __init__(self,data = (), dimension = 0, maxoutlength=2048, wraplength=80, overwrite=True):
53        # print 'Creating new loop block, dimension %d' % dimension
54        self.block = {}
55        self.loops = []
56        self.no_packets = 0
57        self.item_order = []
58        self.lower_keys = []    #for efficiency
59        self.comment_list = {}
60        self.dimension = dimension
61        self.popout = False         #used during load iteration
62        self.curitem = -1           #used during iteration
63        self.maxoutlength = maxoutlength
64        self.wraplength = wraplength
65        self.overwrite = overwrite
66        if not hasattr(self,'loopclass'):  #in case are derived class
67            self.loopclass = LoopBlock  #when making new loops
68        self.char_check = re.compile("[][ \n\r\t!%&\(\)*+,./:<=>?@0-9A-Za-z\\\\^`{}\|~\"#$';_-]+",re.M)
69        if isinstance(data,(TupleType,ListType)):
70            for item in data:
71                self.AddLoopItem(item)
72        elif isinstance(data,LoopBlock):
73            self.block = data.block.copy() 
74            self.item_order = data.item_order[:]
75            self.lower_keys = data.lower_keys[:]
76            self.comment_list = data.comment_list.copy()
77            self.dimension = data.dimension
78            # loops as well; change loop class
79            for loopno in range(len(data.loops)):
80                try:
81                    placeholder = self.item_order.index(data.loops[loopno])
82                except ValueError:
83                    print "Warning: loop %s (%s) in loops, but not in item_order (%s)" % (`data.loops[loopno]`,str(data.loops[loopno]),`self.item_order`)
84                    placeholder = -1
85                self.item_order.remove(data.loops[loopno])   #gone
86                newobject = self.loopclass(data.loops[loopno])
87                # print "Recasting and adding loop %s -> %s" % (`data.loops[loopno]`,`newobject`)
88                self.insert_loop(newobject,position=placeholder)
89
90    def __str__(self):
91        return self.printsection()
92
93    def __setitem__(self,key,value):
94        # catch a one member loop, for convenience
95        # we assume the key is a string value only
96        self.AddLoopItem((key,value))
97
98    def __getitem__(self,key):
99        if isinstance(key,IntType):   #return a packet!!
100            return self.GetPacket(key)       
101        return self.GetLoopItem(key)
102
103    def __delitem__(self,key):
104        self.RemoveLoopItem(key)
105
106    def __len__(self):
107        blen = len(self.block)
108        for aloop in self.loops:
109            # print 'Aloop is %s' % `aloop`
110            blen = blen + len(aloop)  # also a LoopBlock
111        return blen   
112
113    def __nonzero__(self):
114        if self.__len__() > 0: return 1
115        return 0
116
117    # keys returns all internal keys
118    def keys(self):
119        thesekeys = self.block.keys()
120        for aloop in self.loops:
121            thesekeys.extend(aloop.keys())
122        return thesekeys
123
124    def values(self):
125        ourkeys = self.keys()
126        return map(lambda a:self[a],ourkeys)
127
128    def items(self):
129        ourkeys = self.keys()
130        return map(lambda a,b:(a,b),self.keys(),self.values())
131
132    def has_key(self,key):
133        if key.lower() in self.lower_keys:
134            return 1
135        for aloop in self.loops:
136            if aloop.has_key(key): return 1
137        return 0
138
139    def get(self,key,default=None):
140        if self.has_key(key):
141            retval = self.GetLoopItem(key)
142        else:
143            retval = default
144        return retval
145
146    def clear(self):
147        self.block = {}
148        self.loops = []
149        self.item_order = []
150        self.lower_keys = []
151        self.no_packets = 0
152
153    # doesn't appear to work
154    def copy(self):
155        newcopy = self.copy.im_class(dimension = self.dimension)
156        newcopy.block = self.block.copy()
157        newcopy.loops = []
158        newcopy.no_packets = self.no_packets
159        newcopy.item_order = self.item_order[:]
160        newcopy.lower_keys = self.lower_keys[:]
161        for loop in self.loops:
162            try:
163                placeholder = self.item_order.index(loop)
164            except ValueError:
165                print "Warning: loop %s (%s) in loops, but not in item_order (%s)" % (`loop`,str(loop),`self.item_order`)
166                placeholder = -1
167            newcopy.item_order.remove(loop)   #gone
168            newobject = loop.copy()
169            # print "Adding loop %s -> %s" % (`loop`,`newobject`)
170            newcopy.insert_loop(newobject,position=placeholder)
171        return newcopy
172
173    # this is not appropriate for subloops.  Instead, the loop block
174    # should be accessed directly for update
175     
176    def update(self,adict):
177        for key in adict.keys():
178            self.AddLoopItem((key,adict[key]))
179
180    def load_iter(self,coords=[]):
181        count = 0        #to create packet index
182        while not self.popout:
183            # ok, we have a new packet:  append a list to our subloops
184            for aloop in self.loops:
185                aloop.new_enclosing_packet()
186            for iname in self.item_order:
187                if isinstance(iname,LoopBlock):       #into a nested loop
188                    for subitems in iname.load_iter(coords=coords+[count]):
189                        # print 'Yielding %s' % `subitems`
190                        yield subitems
191                    # print 'End of internal loop'
192                else:
193                    if self.dimension == 0:
194                        # print 'Yielding %s' % `self[iname]`
195                        yield self,self[iname]
196                    else:
197                        backval = self.block[iname]
198                        for i in range(len(coords)):
199                           # print 'backval, coords: %s, %s' % (`backval`,`coords`)
200                           backval = backval[coords[i]]
201                        yield self,backval
202            count = count + 1      # count packets
203        self.popout = False        # reinitialise
204        # print 'Finished iterating'
205        yield self,'###Blank###'     #this value should never be used
206
207    # an experimental fast iterator for level-1 loops (ie CIF)
208    def fast_load_iter(self):
209        targets = map(lambda a:self.block[a],self.item_order)
210        while targets:
211            for target in targets:
212                yield self,target
213
214    # Add another list of the required shape to take into account a new outer packet
215    def new_enclosing_packet(self):
216        if self.dimension > 1:      #otherwise have a top-level list
217            for iname in self.keys():  #includes lower levels
218                target_list = self[iname]
219                for i in range(3,self.dimension): #dim 2 upwards are lists of lists of...
220                    target_list = target_list[-1]
221                target_list.append([])
222                # print '%s now %s' % (iname,`self[iname]`)
223
224    def recursive_iter(self,dict_so_far={},coord=[]):
225        # print "Recursive iter: coord %s, keys %s, dim %d" % (`coord`,`self.block.keys()`,self.dimension)
226        my_length = 0
227        top_items = self.block.items()
228        top_values = self.block.values()       #same order as items
229        drill_values = self.block.values()
230        for dimup in range(0,self.dimension):  #look higher in the tree
231            if len(drill_values)>0:            #this block has values
232                drill_values=drill_values[0]   #drill in
233            else:
234                raise StarError("Malformed loop packet %s" % `top_items[0]`)
235        my_length = len(drill_values)
236        if self.dimension == 0:                #top level
237            for aloop in self.loops:
238                for apacket in aloop.recursive_iter():
239                    # print "Recursive yielding %s" % `dict(top_items + apacket.items())`
240                    prep_yield = StarPacket(top_values+apacket.values())  #straight list
241                    for name,value in top_items + apacket.items():
242                        setattr(prep_yield,name,value)
243                    yield prep_yield
244        else:                                  #in some loop
245            for i in range(my_length):
246                kvpairs = map(lambda a:(a,self.coord_to_group(a,coord)[i]),self.block.keys())
247                kvvals = map(lambda a:a[1],kvpairs)   #just values
248                # print "Recursive kvpairs at %d: %s" % (i,`kvpairs`)
249                if self.loops:
250                  for aloop in self.loops:
251                    for apacket in aloop.recursive_iter(coord=coord+[i]):
252                        # print "Recursive yielding %s" % `dict(kvpairs + apacket.items())`
253                        prep_yield = StarPacket(kvvals+apacket.values())
254                        for name,value in kvpairs + apacket.items():
255                            setattr(prep_yield,name,value)
256                        yield prep_yield
257                else:           # we're at the bottom of the tree
258                    # print "Recursive yielding %s" % `dict(kvpairs)`
259                    prep_yield = StarPacket(kvvals)
260                    for name,value in kvpairs:
261                        setattr(prep_yield,name,value)
262                    yield prep_yield
263
264    # small function to use the coordinates.
265    def coord_to_group(self,dataname,coords):
266          if not isinstance(dataname,StringType):
267             return dataname     # flag inner loop processing
268          newm = self[dataname]          # newm must be a list or tuple
269          for c in coords:
270              # print "Coord_to_group: %s ->" % (`newm`),
271              newm = newm[c]
272              # print `newm`
273          return newm
274
275    def flat_iterator(self):
276        if self.dimension == 0:   
277            yield copy.copy(self)
278        else:
279            my_length = 0
280            top_keys = self.block.keys()
281            if len(top_keys)>0:
282                my_length = len(self.block[top_keys[0]])
283            for pack_no in range(my_length):
284                yield(self.collapse(pack_no))
285           
286
287    def insert_loop(self,newloop,position=-1,audit=True):
288        # check that new loop is kosher
289        if newloop.dimension != self.dimension + 1:
290            raise StarError( 'Insertion of loop of wrong nesting level %d, should be %d' % (newloop.dimension, self.dimension+1))
291        self.loops.append(newloop)
292        if audit:
293            dupes = self.audit()
294            if dupes:
295                dupenames = map(lambda a:a[0],dupes)
296                raise StarError( 'Duplicate names: %s' % `dupenames`)
297        if position >= 0:
298            self.item_order.insert(position,newloop)
299        else:
300            self.item_order.append(newloop)
301        # print "Insert loop: item_order now" + `self.item_order`
302
303    def remove_loop(self,oldloop):
304        # print "Removing %s: item_order %s" % (`oldloop`,self.item_order)
305        # print "Length %d" % len(oldloop)
306        self.item_order.remove(oldloop)
307        self.loops.remove(oldloop)
308     
309    def AddComment(self,itemname,comment):
310        self.comment_list[itemname.lower()] = comment
311
312    def RemoveComment(self,itemname):
313        del self.comment_list[itemname.lower()]
314
315    def GetLoopItem(self,itemname):
316        # assume case is correct first
317        try:
318            return self.block[itemname]
319        except KeyError:
320            for loop in self.loops:
321                try:
322                    return loop[itemname]
323                except KeyError:
324                    pass
325        if itemname.lower() not in self.lower_keys:
326            raise KeyError, 'Item %s not in block' % itemname
327        # it is there somewhere, now we need to find it
328        real_keys = self.block.keys()
329        lower_keys = map(lambda a:a.lower(),self.block.keys()) 
330        try:
331            k_index = lower_keys.index(itemname.lower())
332        except ValueError:
333            raise KeyError, 'Item %s not in block' % itemname
334        return self.block[real_keys[k_index]]
335
336    def RemoveLoopItem(self,itemname):
337        if self.has_key(itemname):
338            testkey = itemname.lower()
339            real_keys = self.block.keys()
340            lower_keys = map(lambda a:a.lower(),real_keys)
341            try:
342                k_index = lower_keys.index(testkey)
343            except ValueError:    #must be in a lower loop
344                for aloop in self.loops:
345                    if aloop.has_key(itemname):
346                        # print "Deleting %s (%s)" % (itemname,aloop[itemname])
347                        del aloop[itemname]
348                        if len(aloop)==0:  # all gone
349                           self.remove_loop(aloop)
350                        break
351            else:
352              del self.block[real_keys[k_index]]
353              self.lower_keys.remove(testkey)
354              # now remove the key in the order list
355              for i in range(len(self.item_order)):
356                if isinstance(self.item_order[i],StringType): #may be loop
357                    if self.item_order[i].lower()==testkey:
358                        del self.item_order[i]
359                        break
360            if len(self.block)==0:    #no items in loop, length -> 0
361                self.no_packets = 0
362            return        #no duplicates, no more checking needed
363
364    def AddLoopItem(self,data,precheck=False,maxlength=-1):
365        # print "Received data %s" % `data`
366        # we accept only tuples, strings and lists!!
367        if isinstance(data[0],(TupleType,ListType)):
368           # internal loop
369           # first we remove any occurences of these datanames in
370           # other loops
371           for one_item in data[0]:
372               if self.has_key(one_item):
373                   if not self.overwrite:
374                       raise StarError( 'Attempt to insert duplicate item name %s' % data[0])
375                   else:
376                       del self[one_item]
377           newloop = self.loopclass(dimension = self.dimension+1)
378           keyvals = zip(data[0],data[1])
379           for key,val in keyvals:
380               newloop.AddLoopItem((key,val))
381           self.insert_loop(newloop)
382        elif not isinstance(data[0],StringType):
383                  raise TypeError, 'Star datanames are strings only (got %s)' % `data[0]`
384        else:
385           if data[1] == [] or get_dim(data[1])[0] == self.dimension:
386               if not precheck:
387                   self.check_data_name(data[0],maxlength)    # make sure no nasty characters   
388               # check that we can replace data
389               if not self.overwrite:
390                   if self.has_key(data[0]):
391                       raise StarError( 'Attempt to insert duplicate item name %s' % data[0])
392               # now make sure the data is OK type
393               regval = self.regularise_data(data[1])
394               if not precheck:
395                   try:
396                       self.check_item_value(regval)
397                   except StarError, errmes:
398                       raise StarError( "Item name " + data[0] + " " + `errmes`)
399               if self.dimension > 0:
400                   if self.no_packets <= 0:
401                       self.no_packets = len(data[1])  #first item in this loop
402                   if len(data[1]) != self.no_packets:
403                       raise StarLengthError, 'Not enough values supplied for %s' % (data[0])
404               try:
405                   oldpos = self.GetItemPosition(data[0])
406               except ValueError:
407                   oldpos = len(self.item_order)#end of list
408               self.RemoveLoopItem(data[0])     # may be different case, so have to do this
409               self.block.update({data[0]:regval})  # trust the data is OK
410               self.lower_keys.insert(oldpos,data[0].lower())
411               self.item_order.insert(oldpos,data[0])
412               #    self.lower_keys.append(data[0].lower())
413               #    self.item_order.append(data[0])
414               
415           else:            #dimension mismatch
416               raise StarLengthError, "input data dim %d != required dim %d: %s %s" % (get_dim(data[1])[0],self.dimension,data[0],`data[1]`)
417
418    def check_data_name(self,dataname,maxlength=-1): 
419        if maxlength > 0:
420            if len(dataname)>maxlength:
421                raise StarError( 'Dataname %s exceeds maximum length %d' % (dataname,maxlength))
422        if dataname[0]!='_':
423            raise StarError( 'Dataname ' + dataname + ' does not begin with _')
424        if len (filter (lambda a: ord(a) < 33 or ord(a) > 126, dataname)) > 0:
425            raise StarError( 'Dataname ' + dataname + ' contains forbidden characters')
426
427    def check_item_value(self,item):
428        test_item = item
429        if type(item) != TupleType and type(item) != ListType:
430           test_item = [item]         #single item list
431        def check_one (it):
432            if type(it) == StringType:
433                if it=='': return
434                me = self.char_check.match(it)           
435                if not me:
436                    raise StarError( 'Bad character in %s' % it)
437                else:
438                    if me.span() != (0,len(it)):
439                        raise StarError('Data item "' + it + '"... contains forbidden characters')
440        map(check_one,test_item)
441
442    def regularise_data(self,dataitem):
443        alrighttypes = [IntType, LongType, 
444                        FloatType, StringType]
445        okmappingtypes = [TupleType, ListType]
446        thistype = type(dataitem)
447        if thistype in alrighttypes or thistype in okmappingtypes:
448            return dataitem
449        if isinstance(dataitem,StarTuple) or \
450           isinstance(dataitem,StarList) or \
451           isinstance(dataitem,StarDict):
452            return dataitem
453        # so try to make into a list
454        try:
455            regval = list(dataitem)
456        except TypeError, value:
457            raise StarError( str(dataitem) + ' is wrong type for data value\n' )
458        return regval
459       
460    def GetLoop(self,keyname):
461        if keyname in self.block:        #python 2.2 or above
462            return self
463        for aloop in self.loops:
464            try: 
465                return aloop.GetLoop(keyname)
466            except KeyError:
467                pass
468        raise KeyError, 'Item %s does not exist' % keyname
469
470    def GetPacket(self,index):
471        thispack = StarPacket([])
472        for myitem in self.item_order:
473            if isinstance(myitem,LoopBlock):
474                pack_list = map(lambda b:myitem[b][index],myitem.item_order)
475                # print 'Pack_list -> %s' % `pack_list`
476                thispack.append(pack_list)
477            elif self.dimension==0:
478                thispack.append(self[myitem])
479            else:
480                thispack.append(self[myitem][index])
481                setattr(thispack,myitem,thispack[-1])
482        return thispack
483
484    def AddPacket(self,packet):
485        if self.dimension==0:
486            raise StarError,"Attempt to add packet to top level block"
487        for myitem in self.item_order:
488            self[myitem] = list(self[myitem])   #in case we have stored a tuple
489            self[myitem].append(packet.__getattribute__(myitem))
490        self.no_packets +=1 
491            # print "%s now %s" % (myitem,`self[myitem]`)
492       
493    def RemoveKeyedPacket(self,keyname,keyvalue):
494        packet_coord = list(self[keyname]).index(keyvalue)
495        loophandle = self.GetLoop(keyname)
496        for packet_entry in loophandle.item_order:
497            loophandle[packet_entry] = list(loophandle[packet_entry])
498            del loophandle[packet_entry][packet_coord]
499        self.no_packets -= 1
500       
501    def GetKeyedPacket(self,keyname,keyvalue):
502        #print "Looking for %s in %s" % (keyvalue, self[keyname])
503        one_pack= filter(lambda a:getattr(a,keyname)==keyvalue,self)
504        if len(one_pack)!=1:
505            raise KeyError, "Bad packet key %s = %s: returned %d packets" % (keyname,keyvalue,len(one_pack))
506        #print "Keyed packet: %s" % one_pack[0]
507        return one_pack[0]
508
509    def GetItemOrder(self):
510        return self.item_order[:]
511
512    def ChangeItemOrder(self,itemname,newpos):
513        testpos = self.GetItemPosition(itemname)
514        del self.item_order[testpos]
515        # so we have an object ready for action
516        self.item_order.insert(newpos,itemname)
517
518    def GetItemPosition(self,itemname):
519        import string
520        def low_case(item):
521            try:
522                return string.lower(item)
523            except AttributeError:
524                return item
525        try:
526            testname = string.lower(itemname)
527        except AttributeError: 
528            testname = itemname
529        lowcase_order = map(low_case,self.item_order)
530        return lowcase_order.index(testname)
531
532    def collapse(self,packet_no):
533        if self.dimension == 0:
534            raise StarError( "Attempt to select non-existent packet")
535        newlb = LoopBlock(dimension=self.dimension-1)
536        for one_item in self.item_order:
537            if isinstance(one_item,LoopBlock):
538                newlb.insert_loop(one_item.collapse(packet_no))
539            else:
540                # print "Collapse: %s -> %s" % (one_item,`self[one_item][packet_no]`)
541                newlb[one_item] = self[one_item][packet_no] 
542        return newlb
543       
544    def audit(self):
545        import sets
546        allkeys = self.keys()
547        uniquenames = sets.Set(allkeys)
548        if len(uniquenames) == len(allkeys): return []
549        else:             
550            keycount = map(lambda a:(a,allkeys.count(a)),uniquenames)
551            return filter(lambda a:a[1]>1,keycount)
552       
553    def GetLoopNames(self,keyname):
554        if keyname in self:
555            return self.keys()
556        for aloop in self.loops:
557            try: 
558                return aloop.GetLoopNames(keyname)
559            except KeyError:
560                pass
561        raise KeyError, 'Item does not exist'
562
563    def AddToLoop(self,dataname,loopdata):
564        thisloop = self.GetLoop(dataname)
565        for itemname,itemvalue in loopdata.items():
566            thisloop[itemname] = itemvalue
567
568    def SetOutputLength(self,wraplength=80,maxoutlength=2048):
569        if wraplength > maxoutlength:
570            raise StarError("Wrap length (requested %d) must be <= Maximum line length (requested %d)" % (wraplength,maxoutlength))
571        self.wraplength = wraplength
572        self.maxoutlength = maxoutlength
573        for loop in self.loops:
574            loop.SetOutputLength(wraplength,maxoutlength)
575
576    def printsection(self,instring='',blockstart="",blockend="",indent=0,coord=[]):
577        import cStringIO
578        import string
579        # first make an ordering
580        order = self.item_order[:]
581        # now do it...
582        if not instring:
583            outstring = cStringIO.StringIO()       # the returned string
584        else:
585            outstring = instring
586        if not coord:
587            coords = [0]*(self.dimension-1)
588        else:
589            coords = coord
590        if(len(coords)<self.dimension-1):
591            raise StarError("Not enough block packet coordinates to uniquely define data")
592        # print loop delimiter
593        outstring.write(blockstart)
594        while len(order)>0:
595            # print "Order now: " + `order`
596            itemname = order.pop(0)
597            if self.dimension == 0:            # ie value next to tag
598                if not isinstance(itemname,LoopBlock):  #no loop
599                   # grab any comment
600                   thiscomment = self.comment_list.get(itemname.lower(),'') 
601                   itemvalue = self[itemname]
602                   if isinstance(itemvalue,StringType):  #need to sanitize
603                         thisstring = self._formatstring(itemvalue)
604                   else: thisstring = str(itemvalue)
605                   # try for a tabstop at 40
606                   if len(itemname)<40 and (len(thisstring)-40 < self.wraplength-1):
607                       itemname = itemname + ' '*(40-len(itemname))
608                   else: itemname = itemname + ' '
609                   if len(thisstring) + len(itemname) < (self.wraplength-1):
610                         outstring.write('%s%s' % (itemname,thisstring))
611                         if thiscomment:
612                             if len(thiscomment)+len(thisstring)+len(itemname)< (self.wraplength-3):
613                                 outstring.write(' #'+thiscomment)
614                   else:
615                         outstring.write('%s\n %s' % (itemname, thisstring))
616                         if thiscomment:
617                             if len(thiscomment)+len(thisstring)<(self.wraplength-3):
618                                 outstring.write(' #'+thiscomment)
619                             else:
620                                 outstring.write('\n#'+thiscomment)
621                   outstring.write('\n')
622                else:   # we are asked to print an internal loop block
623                    #first make sure we have sensible coords.  Length should be one
624                    #less than the current dimension
625                    outstring.write(' '*indent); outstring.write('loop_\n')
626                    itemname.format_names(outstring,indent+2)
627                    itemname.format_packets(outstring,coords,indent+2)
628            else:   # we are a nested loop
629                outstring.write(' '*indent); outstring.write('loop_\n')
630                self.format_names(outstring,indent+2)
631                self.format_packets(outstring,coords,indent+2)
632        if instring: return   #inside a recursion
633        else:
634            returnstring = outstring.getvalue()
635        outstring.close()
636        return returnstring
637
638    def format_names(self,outstring,indent=0):
639        temp_order = self.item_order[:]
640        while len(temp_order)>0:
641            itemname = temp_order.pop(0)
642            if isinstance(itemname,StringType):  #(not loop)
643                outstring.write(' ' * indent) 
644                outstring.write(itemname)
645                outstring.write("\n")
646            else:                                # a loop
647                outstring.write(' ' * indent) 
648                outstring.write("loop_\n")
649                itemname.format_names(outstring,indent+2)
650                outstring.write(" stop_\n")
651
652    def format_packets(self,outstring,coordinates,indent=0):
653       import cStringIO
654       import string
655       # get our current group of data
656       # print 'Coords: %s' % `coordinates`
657       alldata = map(lambda a:self.coord_to_group(a,coordinates),self.item_order)
658       # print 'Alldata: %s' % `alldata`
659       packet_data = apply(zip,alldata)
660       # print 'Packet data: %s' % `packet_data`
661       curstring = ''
662       for position in range(len(packet_data)):
663           for point in range(len(packet_data[position])):
664               datapoint = packet_data[position][point]
665               packstring = self.format_packet_item(datapoint,indent)
666               if len(curstring) + len(packstring)> self.wraplength-2: #past end of line with space
667                   curstring = curstring + '\n' + ' '*indent + packstring
668               elif curstring == '':
669                   curstring = curstring + ' '*indent + packstring
670               else:
671                   curstring = curstring + ' ' + packstring
672           outstring.write(curstring + '\n')     #end of one packet
673           curstring = ''
674       outstring.write(' ' + curstring + '\n')    #last time through
675               
676    def format_packet_item(self,pack_item,indent):
677        # print 'Formatting %s' % `pack_item`
678        curstring = ''
679        if isinstance(pack_item,(StringType,IntType,FloatType,LongType,StarTuple,StarList)):
680           if isinstance(pack_item,StringType):
681               thisstring = self._formatstring(pack_item) #no spaces yet
682               if '\n' in thisstring:    #must have semicolon digraph then
683                   curstring = curstring + thisstring
684                   curstring = curstring + (' ' * indent)
685                   thisstring = ''
686           else: 
687               thisstring = '%s' % str(pack_item)
688           if len(curstring) + len(thisstring)> self.wraplength-2: #past end of line with space
689               curstring = curstring + '\n' #add the space
690               curstring = curstring + (' ' * indent) + thisstring
691           else: 
692               curstring = curstring + ' ' + thisstring
693        # Now, for each nested loop we call ourselves again
694        # After first outputting the current line
695        else:               # a nested packet
696           if not isinstance(pack_item[0],(ListType,TupleType)):  #base packet
697               item_list = pack_item
698           else:
699               item_list = apply(zip,pack_item)
700           for sub_item in item_list:
701               curstring = curstring + ' ' + self.format_packet_item(sub_item,indent)
702           # stop_ is not issued at the end of each innermost packet
703           if isinstance(pack_item[0],(ListType,TupleType)):
704               curstring = curstring + ' stop_ '
705        return curstring         
706
707    def _formatstring(self,instring):
708        import string
709        if len(instring)==0: return "''"
710        if len(instring)< (self.maxoutlength-2) and '\n' not in instring and not ('"' in instring and '\'' in instring):
711            if not ' ' in instring and not '\t' in instring and not '\v' \
712              in instring and not '_' in instring and not (instring[0]=="'" or \
713                 instring[0]=='"'):                  # no blanks
714                return instring
715            if not "'" in instring:                                       #use apostrophes
716                return "'%s'" % (instring)
717            elif not "\"" in instring:
718                return '"%s"' % (instring)
719        # is a long one or one that needs semicolons due to carriage returns
720        outstring = "\n;"
721        # if there are returns in the string, try to work with them
722        while 1:
723            retin = string.find(instring,'\n')+1
724            if retin < self.maxoutlength and retin > 0:      # honour this break
725                outstring = outstring + instring[:retin]
726                instring = instring[retin:]
727            elif len(instring)<self.maxoutlength:            # finished
728                outstring = outstring + instring + '\n;\n'
729                break
730            else:                             # find a space
731                for letter in range(self.maxoutlength-1,self.wraplength-1,-1): 
732                    if instring[letter] in ' \t\f': break
733                outstring = outstring + instring[:letter+1]
734                outstring = outstring + '\n'
735                instring = instring[letter+1:]           
736        return outstring
737
738
739
740class StarBlock(LoopBlock):
741    def __init__(self,*pos_args,**keyword_args):
742        LoopBlock.__init__(self,*pos_args,**keyword_args)
743        self.saves = BlockCollection(element_class=LoopBlock,type_tag="save")
744
745    def __getitem__(self,key):
746        if key == "saves":
747            return self.saves
748        else:
749            return LoopBlock.__getitem__(self,key)
750
751    def __setitem__(self,key,value):
752        if key == "saves":
753            self.saves[key] = value
754        else:
755            LoopBlock.__setitem__(self,key,value)
756
757    def clear(self):
758        LoopBlock.clear(self)
759        self.saves = BlockCollection(element_class=LoopBlock,type_tag="save_")
760
761    def copy(self):
762        newblock = LoopBlock.copy(self)
763        newblock.saves = self.saves.copy()
764        return self.copy.im_class(newblock)   #catch inheritance
765
766    def has_key(self,key):
767        if key == "saves": return 1
768        else: return LoopBlock.has_key(self,key)
769       
770    def __str__(self):
771        retstr = ''
772        for sb in self.saves.keys(): 
773            retstr = retstr + '\nsave_%s\n\n' % sb
774            self.saves[sb].SetOutputLength(self.wraplength,self.maxoutlength)
775            retstr = retstr + str(self.saves[sb])
776            retstr = retstr + '\nsave_\n\n'
777        return retstr + LoopBlock.__str__(self)
778
779
780class StarPacket(list):
781    pass
782
783class BlockCollection:
784    def __init__(self,datasource=None,element_class=StarBlock,type_tag=''):
785        self.dictionary = {}
786        self.type_tag = type_tag
787        self.lower_keys = []              # for efficiency
788        self.element_class = element_class
789        if isinstance(datasource,(DictType,BlockCollection)):
790            for key,value in datasource.items():
791                if value.__class__ == element_class:
792                    self[key]=value
793                else:
794                    self[key]= element_class(value)
795        self.header_comment = ''
796     
797    def __str__(self):
798        return self.WriteOut()
799
800    def __setitem__(self,key,value):
801        if isinstance(value,(self.element_class,DictType)):
802            self.NewBlock(key,value,replace=True)
803        else: raise TypeError
804        self.lower_keys.append(key.lower())
805
806    # due to attempt to get upper/lower case treated as identical
807    # we have a bit of cruft here
808    def __getitem__(self,key):
809        try:
810            return self.dictionary[key]
811        except KeyError:
812            if key.lower() not in self.lower_keys:
813                raise KeyError, "No such item: %s" % key
814        curr_keys = self.dictionary.keys()
815        lower_ordered = map(lambda a:a.lower(),curr_keys)
816        keyindex = lower_ordered.index(key.lower())
817        return self.dictionary[curr_keys[keyindex]]
818
819    # we have to get an ordered list of the current keys,
820    # as we'll have to delete one of them anyway
821    def __delitem__(self,key):
822        try:
823            del self.dictionary[key]
824            self.lower_keys.remove(key.lower())
825        except KeyError:
826            if not self.has_key(key):
827                raise KeyError
828            curr_keys = self.dictionary.keys()
829            lower_ordered = map(lambda a:a.lower(),curr_keys)
830            keyindex = lower_ordered.index(key.lower())
831            del self.dictionary[curr_keys[keyindex]]
832       
833    def __len__(self):
834        return len(self.dictionary)
835
836    def keys(self):
837        return self.dictionary.keys()
838
839    # changes to take case independence into account
840    def has_key(self,key):
841        if not isinstance(key,StringType): return 0
842        if self.dictionary.has_key(key):
843           return 1
844        if key.lower() in self.lower_keys:
845           return 1
846        return 0
847
848    def get(self,key,default=None):
849        if self.dictionary.has_key(key):
850            return self.dictionary[key]
851        elif self.has_key(key):     # take account of case
852            return self.__getitem__(key)
853        else:
854            return default
855
856    def clear(self):
857        self.dictionary.clear()
858        self.lower_keys = []
859
860    def copy(self):   
861        newcopy = self.dictionary.copy()
862        return BlockCollection('',newcopy)
863     
864    def update(self,adict):
865        for key in adict.keys():
866            self.dictionary[key] = adict[key]
867        self.lower_keys.extend(map(lambda a:a.lower(),adict.keys()))
868
869    def items(self):
870        return self.dictionary.items()
871
872    def first_block(self):
873        if self.keys():
874            return self[self.keys()[0]]
875
876    def NewBlock(self,blockname,blockcontents=(),replace=False,fix=True):
877        if not blockcontents:
878            blockcontents = self.element_class()
879        elif isinstance(blockcontents,DictType):
880            blockcontents = self.element_class(blockcontents)
881        if not isinstance(blockcontents,self.element_class):
882            raise StarError( 'Block is not of required type %s, is %s' % self.element_class.__name__,blockcontents.__class__.__name__)
883        if fix:
884            newblockname = re.sub('[  \t]','_',blockname)
885        else: newblockname = blockname
886        new_lowerbn = newblockname.lower()
887        if self.lower_keys.count(new_lowerbn):    #already in CIF
888            if not replace:
889                raise StarError( "Attempt to replace existing block" + blockname)
890            # generate a list of lower-case keys in correct order
891            current_keys = self.dictionary.keys()
892            blocknames = map(lambda a:a.lower(),current_keys)
893            location = blocknames.index(new_lowerbn)
894            del self.dictionary[current_keys[location]]
895            self.lower_keys.remove(new_lowerbn)
896        self.dictionary.update({blockname:blockcontents})
897        self.lower_keys.append(new_lowerbn)
898
899    def merge(self,new_bc,mode="strict",single_block=[],
900                   idblock="",match_att=[],match_function=None):
901        if single_block:
902            self.dictionary[single_block[0]].merge(new_bc[single_block[1]],mode,
903                                                   match_att=match_att,
904                                                   match_function=match_function)
905            return None
906        base_keys = self.keys()
907        block_to_item = base_keys   #default
908        new_keys = new_bc.keys()
909        if match_att:
910            #make a blockname -> item name map
911            if match_function:
912                block_to_item = map(lambda a:match_function(self[a]),self.keys())
913            else:
914                block_to_item = map(lambda a:self[a].get(match_att[0],None),self.keys())
915            #print `block_to_item`
916        for key in new_keys:
917            if key == idblock: continue
918            basekey = key        #default value
919            attval = new_bc[key].get(match_att[0],0)
920            for ii in range(len(block_to_item)):  #do this way to get looped names
921                thisatt = block_to_item[ii]
922                #print "Looking for %s in %s" % (attval,thisatt)
923                if attval == thisatt or \
924                   (isinstance(thisatt,ListType) and attval in thisatt):
925                      basekey = base_keys.pop(ii)
926                      block_to_item.remove(thisatt)
927                      break
928            if not self.dictionary.has_key(basekey) or mode=="replace":
929                self.dictionary[basekey] = new_bc[key]
930            else:
931                if mode=="strict":
932                    raise StarError( "In strict merge mode: block %s in old and block %s in new files" % (basekey,key))
933                elif mode=="overlay":
934                    # print "Merging block %s with %s" % (basekey,key)
935                    self.dictionary[basekey].merge(new_bc[key],mode,match_att=match_att)
936                else: 
937                    raise StarError( "Merge called with unknown mode %s" % mode)
938
939    def get_all(self,item_name):
940        raw_values = map(lambda a:self[a].get(item_name),self.dictionary.keys())
941        raw_values = filter(lambda a:a != None, raw_values)
942        ret_vals = []
943        for rv in raw_values:
944            if isinstance(rv,ListType):
945                for rvv in rv:
946                    if rvv not in ret_vals: ret_vals.append(rvv)
947            else:
948                if rv not in ret_vals: ret_vals.append(rv)
949        return ret_vals
950
951    def WriteOut(self,comment='',wraplength=80,maxoutlength=2048):
952        import cStringIO
953        if not comment:
954            comment = self.header_comment
955        outstring = cStringIO.StringIO()
956        outstring.write(comment)
957        for datablock in self.dictionary.keys():
958            outstring.write('\n' + self.type_tag +datablock+'\n')
959            self.dictionary[datablock].SetOutputLength(wraplength,maxoutlength)
960            outstring.write(str(self.dictionary[datablock]))
961        returnstring =  outstring.getvalue()
962        outstring.close()
963        return returnstring
964
965
966class StarFile(BlockCollection):
967    def __init__(self,datasource=None,maxinlength=-1,maxoutlength=0,blocktype=StarBlock,**kwargs):
968        BlockCollection.__init__(self,datasource=datasource,element_class=blocktype,type_tag='data_')
969        if isinstance(datasource, StarFile):
970            self.my_uri = datasource.my_uri
971        self.maxinlength = maxinlength      #no restriction
972        if maxoutlength == 0:
973            self.maxoutlength = 2048 
974        else:
975            self.maxoutlength = maxoutlength
976        if type(datasource) is StringType or hasattr(datasource,"read"):
977            newself = ReadStar(datasource,self.maxinlength,**kwargs)
978            # print "Reinjecting by calling %s.__init__ with kwargs %s" % (`self.__init__.im_class`,kwargs)
979            self.__init__.im_class.__init__(self,datasource=newself,maxoutlength=maxoutlength,**kwargs)
980        self.header_comment = \
981"""#\\#STAR
982##########################################################################
983#               STAR Format file
984#               Produced by PySTARRW module
985#
986#  This is a STAR file.  STAR is a superset of the CIF file type.  For
987#  more information, please refer to International Tables for Crystallography,
988#  Volume G, Chapter 2.1
989#
990##########################################################################
991"""
992    def set_uri(self,my_uri): self.my_uri = my_uri
993
994
995class StarError(Exception):
996    def __init__(self,value):
997        self.value = value
998    def __str__(self):
999        return '\nStar Format error: '+ self.value
1000
1001class StarLengthError(Exception):
1002    def __init__(self,value):
1003        self.value = value
1004    def __str__(self):
1005        return '\nStar length error: ' + self.value
1006def ReadStar(filename,maxlength=2048,dest=StarFile(),scantype='standard',grammar='1.1'):
1007    import string
1008    if grammar=="1.1":
1009        import YappsStarParser_1_1 as Y
1010    elif grammar=="1.0":
1011        import YappsStarParser_1_0 as Y
1012    elif grammar=="DDLm":
1013        import YappsStarParser_DDLm as Y
1014    if isinstance(filename,basestring):
1015        filestream = urlopen(filename)
1016    else:
1017        filestream = filename   #already opened for us
1018    my_uri = ""
1019    if hasattr(filestream,"geturl"): 
1020        my_uri = filestream.geturl()
1021    text = filestream.read()
1022    if isinstance(filename,basestring): #we opened it, we close it
1023        filestream.close()
1024    if not text:      # empty file, return empty block
1025        dest.set_uri(my_uri)
1026        return dest
1027    # we recognise ctrl-Z as end of file
1028    endoffile = text.find('\x1a')
1029    if endoffile >= 0: 
1030        text = text[:endoffile]
1031    split = string.split(text,'\n')
1032    if maxlength > 0:
1033        toolong = filter(lambda a:len(a)>maxlength,split)
1034        if toolong:
1035            pos = split.index(toolong[0])
1036            raise StarError( 'Line %d contains more than %d characters' % (pos+1,maxlength))
1037    try: 
1038        if scantype == 'standard':
1039            parser = Y.StarParser(Y.StarParserScanner(text))
1040        else:
1041            parser = Y.StarParser(Y.yappsrt.Scanner(None,[],text,scantype='flex'))
1042        proto_star = getattr(parser,"input")()
1043    except Y.yappsrt.SyntaxError:
1044        errorstring = 'Syntax error in input file: last value parsed was %s' % Y.lastval
1045        errorstring = errorstring + '\nParser status: %s' % `parser._scanner`
1046        raise StarError( errorstring)
1047    # duplication check on all blocks
1048    audit_result = map(lambda a:(a,proto_star[a].audit()),proto_star.keys())
1049    audit_result = filter(lambda a:len(a[1])>0,audit_result)
1050    if audit_result:
1051        raise StarError( 'Duplicate keys as follows: %s' % `audit_result`)
1052    proto_star.set_uri(my_uri)
1053    return proto_star
1054
1055def get_dim(dataitem,current=0,packlen=0):
1056    zerotypes = [IntType, LongType, 
1057                    FloatType, StringType]
1058    if type(dataitem) in zerotypes:
1059        return current, packlen
1060    if not dataitem.__class__ == ().__class__ and \
1061       not dataitem.__class__ == [].__class__:
1062       return current, packlen
1063    elif len(dataitem)>0: 
1064    #    print "Get_dim: %d: %s" % (current,`dataitem`)
1065        return get_dim(dataitem[0],current+1,len(dataitem))
1066    else: return current+1,0
1067   
1068
1069
Note: See TracBrowser for help on using the repository browser.