1# -*- test-case-name: foolscap.test.test_banana -*- 2 3import six 4from twisted.internet.defer import Deferred 5from foolscap.constraint import Any, ByteStringConstraint 6from foolscap.tokens import Violation, BananaError, INT, STRING 7from foolscap.slicer import BaseSlicer, BaseUnslicer, LeafUnslicer 8from foolscap.slicer import BananaUnslicerRegistry 9 10class ReplaceVocabularyTable: 11 pass 12 13class AddToVocabularyTable: 14 pass 15 16class ReplaceVocabSlicer(BaseSlicer): 17 # this works somewhat like a dictionary 18 opentype = ('set-vocab',) 19 trackReferences = False 20 21 def slice(self, streamable, banana): 22 # we need to implement slice() (instead of merely sliceBody) so we 23 # can get control at the beginning and end of serialization. It also 24 # gives us access to the Banana protocol object, so we can manipulate 25 # their outgoingVocabulary table. 26 self.streamable = streamable 27 self.start(banana) 28 for o in self.opentype: 29 yield six.ensure_binary(o) 30 # the vocabDict maps strings to index numbers. The far end needs the 31 # opposite mapping, from index numbers to strings. We perform the 32 # flip here at the sending end. 33 stringToIndex = self.obj 34 for s in stringToIndex.keys(): 35 assert isinstance(s, six.binary_type), "%r %s" % (s, type(s)) 36 indexToString = dict([(stringToIndex[s],s) for s in stringToIndex]) 37 assert len(stringToIndex) == len(indexToString) # catch duplicates 38 indices = list(indexToString.keys()) 39 indices.sort() 40 for index in indices: 41 string = indexToString[index] 42 yield index 43 yield six.ensure_binary(string) 44 self.finish(banana) 45 46 def start(self, banana): 47 # this marks the transition point between the old vocabulary dict and 48 # the new one, so now is the time we should empty the dict. 49 banana.outgoingVocabTableWasReplaced({}) 50 51 def finish(self, banana): 52 # now we replace the vocab dict 53 banana.outgoingVocabTableWasReplaced(self.obj) 54 55class ReplaceVocabUnslicer(LeafUnslicer): 56 """Much like DictUnslicer, but keys must be numbers, and values must be 57 strings. This is used to set the entire vocab table at once. To add 58 individual tokens, use AddVocabUnslicer by sending an (add-vocab num 59 string) sequence.""" 60 opentype = ('set-vocab',) 61 unslicerRegistry = BananaUnslicerRegistry 62 maxKeys = None 63 valueConstraint = ByteStringConstraint(100) 64 65 def setConstraint(self, constraint): 66 if isinstance(constraint, Any): 67 return 68 assert isinstance(constraint, ByteStringConstraint) 69 self.valueConstraint = constraint 70 71 def start(self, count): 72 self.d = {} 73 self.key = None 74 75 def checkToken(self, typebyte, size): 76 if self.maxKeys is not None and len(self.d) >= self.maxKeys: 77 raise Violation("the table is full") 78 if self.key is None: 79 if typebyte != INT: 80 raise BananaError("VocabUnslicer only accepts INT keys") 81 else: 82 if typebyte != STRING: 83 raise BananaError("VocabUnslicer only accepts STRING values") 84 if self.valueConstraint: 85 self.valueConstraint.checkToken(typebyte, size) 86 87 def receiveChild(self, token, ready_deferred=None): 88 assert not isinstance(token, Deferred) 89 assert ready_deferred is None 90 if self.key is None: 91 if token in self.d: 92 raise BananaError("duplicate key '%s'" % token) 93 self.key = token 94 else: 95 self.d[self.key] = token 96 self.key = None 97 98 def receiveClose(self): 99 if self.key is not None: 100 raise BananaError("sequence ended early: got key but not value") 101 # now is the time we replace our protocol's vocab table 102 self.protocol.replaceIncomingVocabulary(self.d) 103 return ReplaceVocabularyTable, None 104 105 def describe(self): 106 if self.key is not None: 107 return "<vocabdict>[%s]" % self.key 108 else: 109 return "<vocabdict>" 110 111 112class AddVocabSlicer(BaseSlicer): 113 opentype = ('add-vocab',) 114 trackReferences = False 115 116 def __init__(self, value): 117 assert isinstance(value, str) 118 self.value = value 119 120 def slice(self, streamable, banana): 121 # we need to implement slice() (instead of merely sliceBody) so we 122 # can get control at the beginning and end of serialization. It also 123 # gives us access to the Banana protocol object, so we can manipulate 124 # their outgoingVocabulary table. 125 self.streamable = streamable 126 self.start(banana) 127 for o in self.opentype: 128 yield six.ensure_binary(o) 129 yield self.index 130 yield self.value 131 self.finish(banana) 132 133 def start(self, banana): 134 # this marks the transition point between the old vocabulary dict and 135 # the new one, so now is the time we should decide upon the key. It 136 # is important that we *do not* add it to the dict yet, otherwise 137 # we'll send (add-vocab NN [VOCAB#NN]), which is kind of pointless. 138 index = banana.allocateEntryInOutgoingVocabTable(self.value) 139 self.index = index 140 141 def finish(self, banana): 142 banana.outgoingVocabTableWasAmended(self.index, self.value) 143 144class AddVocabUnslicer(BaseUnslicer): 145 # (add-vocab num string): self.vocab[num] = string 146 opentype = ('add-vocab',) 147 unslicerRegistry = BananaUnslicerRegistry 148 index = None 149 value = None 150 valueConstraint = ByteStringConstraint(100) 151 152 def setConstraint(self, constraint): 153 if isinstance(constraint, Any): 154 return 155 assert isinstance(constraint, ByteStringConstraint) 156 self.valueConstraint = constraint 157 158 def checkToken(self, typebyte, size): 159 if self.index is None: 160 if typebyte != INT: 161 raise BananaError("Vocab key must be an INT") 162 elif self.value is None: 163 if typebyte != STRING: 164 raise BananaError("Vocab value must be a STRING") 165 if self.valueConstraint: 166 self.valueConstraint.checkToken(typebyte, size) 167 else: 168 raise Violation("add-vocab only accepts two values") 169 170 def receiveChild(self, obj, ready_deferred=None): 171 assert not isinstance(obj, Deferred) 172 assert ready_deferred is None 173 if self.index is None: 174 self.index = obj 175 else: 176 self.value = obj 177 178 def receiveClose(self): 179 if self.index is None or self.value is None: 180 raise BananaError("sequence ended too early") 181 self.protocol.addIncomingVocabulary(self.index, self.value) 182 return AddToVocabularyTable, None 183 184 def describe(self): 185 if self.index is not None: 186 return "<add-vocab>[%d]" % self.index 187 return "<add-vocab>" 188