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