<commit_before>#!/usr/bin/python
# Toy Reduced Awful Programming Language
# Copyright (c) 2014 Alexander Sosedkin <monk@unboiled.info>, see LICENSE file

# This program is an interpreter for a toy language with simple syntax rules:
# (try reading the examples in test.py, most are pretty self-explanatory!)
# 1) Braces define priority
# 2) Applications chains: 'obj1 obj2 obj3' means obj3 gets applied to
#    the result of obj2 to obj1 application and so on
# Application usually means a dictionary lookup in first object,
# but there also are other important effects of application (see trapl_apply):
# 1) if first is a 'callable' it calls a wrapped python functions of 1 argument
# 2) if second is a 'method' it gets applied to the first object
# Additional syntax is supported with source-to-source transformations.
# The interpreter reuses python stack, closures... staying relatively compact.
# For even simpler, but less capable versions, please refer to the previous
# versions in the VCS (first version was very simple and < 87 lines long!).

class TRAPLError(RuntimeError): pass # a base exception for interpreter errors

class mydict(dict): # a lousily implemented immutable dictionary
    def __getattr__(self, k):    return self[k]
    def __setitem__(self, k, v): raise TRAPLError('__setitem__ disabled')
    def __setattr__(self, k, v): raise TRAPLError('__setattr__ disabled')
    def clear(self, *a, **kwar): raise TRAPLError('clear disabled')
    def update(self, *a, **kwa): raise TRAPLError('update disabled')
    def setdefault(s, *a, **kw): raise TRAPLError('setdefault disabled')
    def __call__(self, *a, **k): # this is how you extend mydict:
        n = self.copy()          # mydict({'a': 'b'})({'c': 'd'})(x='y')
        n.update(dict(*a, **k))  # results in a new mydict equivalent to
        return mydict(n)         # mydict({'a': 'b', 'c': 'd', 'x': 'y'})

OBJ = mydict() # create an immutable object, the base for all future objects
# methods and callables are objects with special fields (see trapl_apply)
METH = lambda body, **kwa: OBJ(_meth_=body, **kwa) # method and callables are
CALL = lambda body, **kwa: OBJ(_call_=body, **kwa) # created like this
# this is how you create a method of obj a that takes an additional parameter:
EQ = METH(lambda a: CALL(lambda b: # this means the method returns a callable
    TRAPL['true'] if a._val_ == b._val_ else TRAPL['false'] # which returns this
))
TRAPL = OBJ({ # let's define the standard library, accessible later as 'trapl'
    'obj': OBJ, # provide a means of accessing the original, blank object
    'false': OBJ({ # an instance of a boolean value false
        '_val_': False, # put a python-object value in _val_
        'not': METH(lambda b: b(_val_=(not b._val_))), # and several methods
        # understanding how they work remains as an exercise
        'or': METH(lambda b: CALL(lambda arg:
            arg if not b._val_ else b # this way 'false or 2' is 2
        )),
        'and': METH(lambda b: CALL(lambda arg:
            arg if b._val_ else b # this way 'true and 2' is 2
        )),
        'eq': EQ, # this is the example universal method from the above
    }),
    'int': OBJ({
        '_val_': 0,
        'new': METH(lambda i: CALL(lambda s: i(_val_=int(s._val_)))),
        'neg': METH(lambda i: i(_val_=-i._val_)),
        'inc': METH(lambda i: i(_val_=(i._val_ + 1))),
        'add': METH(lambda a: CALL(lambda b: a(_val_=a._val_ + b._val_))),
        'str': METH(lambda a: TRAPL['str'](_val_=str(a._val_))),
        'ge': METH(lambda a: CALL(lambda b:
            TRAPL['true' if a._val_ >= b._val_ else 'false']
        )),
        'eq': EQ,
    }),
    'str': OBJ({
        '_val_': '',
        'new': METH(lambda s: CALL(lambda n: s(_val_=n._val_))),
        'cat': METH(lambda s: CALL(lambda n: s(_val_=(s._val_ + n._val_)))),
        'len': METH(lambda s: TRAPL['int'](_val_=len(s._val_))),
        'rev': METH(lambda s: s(_val_=s._val_[::-1])),
        'dec': METH(lambda s: CALL(lambda e: s(_val_=decode_str(e._val_)))),
        'eq': EQ,
    }),
    'ext': CALL(lambda o: CALL(lambda n: CALL(lambda w: o({n._val_: w})))),
    'code': OBJ(_magic_='code'), # objects with _magic_ are special
    'eval': CALL(lambda code: OBJ(_magic_='eval', _magic_code_=code)),
    'drop': CALL(lambda name: OBJ(_magic_='drop', _magic_name_=name._val_)),
    'with': CALL(lambda name: CALL(lambda val:
        OBJ(_magic_='with', _magic_name_=name._val_, _magic_value_=val)
    )), # injects {name._val_: val} in current context, see trapl_eval
    'get': CALL(lambda name: OBJ(_magic_='get', _magic_name_=name._val_)),
    'func': CALL(lambda arg_name: CALL(lambda code:
        OBJ(_magic_='func', _magic_code_=code, _magic_arg_name_=arg_name),
    )), # make a callable will evaluate code with {arg_name._val_: arg_val}
    'atch': CALL(lambda o: CALL(lambda name: CALL(lambda call:
            o({name._val_: METH(call._call_)})
    ))), # Attaches a callable call to an object o as amethod with name name
    'dtch': CALL(lambda o: CALL(lambda name:
        CALL(o[name._val_]._meth_)
    )), # Returns a method detached from an object and turned into a callable
    'if': CALL(lambda cond:
        CALL(lambda val_true: CALL(lambda unused_val_false: val_true))
        if cond._val_ == True else
        CALL(lambda unused_val_true: CALL(lambda val_false: val_false))
    ),
})

def parse(tokens): # turns 'a ( b ( c ) d )' into ['a', ['b', ['c'], 'd']]
    if isinstance(tokens, str): tokens = tokens.split()
    tree = []
    while tokens:
        t = tokens.pop(0)
        if t == '(': tree.append(parse(tokens))
        elif t == ')': break
        else: tree.append(t)
    return tree

# Does the reverse job
flatten = lambda t: \
    ' '.join('( ' + flatten(a) + ' )' if isinstance(a, list) else a for a in t)

def trapl_apply(to, what): # Handles how one object gets applied to another
    if '_call_' in to: return to._call_(what) # appying to a callable calls it
    if not '_val_' in what: raise TRAPLError('No _val_ in %s' % what)
    if what._val_ in to: # if it's a string mathcing one of the object's keys
        x = to[what._val_] # retrieve it from the object
    # TODO LATER: Implement this functionality in trapl, remove from here
    elif '_get_' in to: # otherwise call a special catchall callable if present
        x = trapl_apply(trapl_apply(to, TRAPL['str'](_val_='_get_')), what)
    else:
        raise TRAPLError('No %s in %s' % (what._val_, to.keys()))
    if '_meth_' in x: return x._meth_(to) # if the result is a method, apply it
    # Note: to get a method from the object without appying it, see dtach
    return x

def _trapl_eval(tree, context=None, default_object=OBJ): # evaluate a tree
    context = {'trapl': TRAPL} if context is None else context.copy()
    curr = None
    while tree:
        next, tree = tree[0], tree[1:]
        if isinstance(next, list): next = _trapl_eval(next, context)
        if isinstance(next, str): # Autocast unknown literals to strings
            next = context.get(next, TRAPL['str'](_val_=next))
        curr = trapl_apply(curr, next) if not curr is None else next
        while '_magic_' in curr: # Handle special messages to the interpreter
            if curr._magic_ == 'with': # inject a value in current context
                context[curr._magic_name_] = curr._magic_value_
                curr = TRAPL['ign']
            elif curr._magic_ == 'drop': # inject a value in current context
                del context[curr._magic_name_]
                curr = TRAPL['ign']
            elif curr._magic_ == 'get': # get a value from current context
                curr = context[curr._magic_name_] # see '4+' example in test.py
            elif curr._magic_ == 'eval': # evaluate a string
                utree = parse(curr._magic_code_._val_)
                ctx = context.copy()
                if '_magic_definition_context_' in curr:
                    ctx.update(curr._magic_definition_context_)
                if '_magic_arg_val_' in curr:
                    ctx[curr._magic_arg_name_._val_] = curr._magic_arg_val_
                curr = _trapl_eval(utree, ctx)
            elif curr._magic_ == 'func': # create a function (closure)
                ctx, arg_name = context.copy(), curr._magic_arg_name_
                code = curr._magic_code_
                curr = CALL(lambda arg_val:
                    OBJ(_magic_='eval',
                        _magic_code_=code,
                        _magic_arg_name_=arg_name,
                        _magic_arg_val_=arg_val,
                        _magic_definition_context_=context,
                ))
            elif curr._magic_ == 'code': # create a string from remaining code
                # fall out of current brace if empty (allows trapl.code)
                if not tree: return curr
                curr = TRAPL['str'](_val_=flatten(tree))
                tree = []
    return curr or default_object

syntax_plain = lambda code: code # source code transformations may be used
# Let's define some code transformations to make syntax more pleasant
import re, base64
encode_str = lambda s: 'ENC' + base64.b32encode(s).replace('=', '0')
decode_str = lambda s: base64.b32decode(s[3:].replace('0', '='))
include_str = lambda s: ' ( trapl str dec %s ) ' % encode_str(s)
quotes = lambda code: re.sub(r"\B'([^'\\]*(?:\\.[^'\\]*)*)'\B", lambda m:
    include_str(m.group(1).replace('\\\'', '\'').replace('\\\\', '\\')), code)
dots = lambda code: re.sub(r"(\w*(?:\.(?:\w*))+)", lambda m:
    ' ( ' + m.group(1).split('.')[0] + ' ' +
    ' '.join(include_str(s) for s in m.group(1).split('.')[1:]) + ' ) ', code)
assign = lambda code: re.sub(r"(\w*)(?:\s*)=", lambda m:
    'trapl with %s ' % include_str(m.group(1)), code)

curly_func = lambda code: flatten(_curly_func(parse(
    code.replace('{', ' ( { ').replace('}', ' } ) ').replace('|', ' | ')
)))
funcize = lambda argname, tree: \
    ['trapl', 'func', parse(include_str(argname)), ['trapl', 'code'] + tree]
def _curly_func(tree):
    if isinstance(tree, str): return tree
    if len(tree) > 3:
        if tree[0] == '{' and '|' in tree[1:-1] and tree[-1] == '}':
            i = tree[1:-1].index('|') + 1
            args, tree = tree[1:i], tree[i+1:-1]
            for a in args[::-1]:
                tree = funcize(a, tree)
    return [_curly_func(t) for t in tree]

autoint = lambda code: re.sub(r"\b(\d+)\b", lambda m:
    ' ( trapl int new ' + include_str(m.group(1)) + ' ) ', code)

def syntax_rich(code): # apply lots of source-to-source transformations
    # convert 'a' into ( trapl string dec ENCSMTH ) to protect it
    code = quotes(code) # from futher damage, respects escaping
    # convert {a b|b a} -> (trapl func a ( trapl code ( trapl func b (
    code = code.replace('(', ' ( ').replace(')', ' ) ') # HACK: deduplicate
    code = curly_func(code) # trapl code b a ) ) ) )
    code = autoint(code) # convert 3 -> ( trapl int new 3 )
    code = dots(code) # convert t = trapl.true -> t = (trapl true)
    code = assign(code) # convert t=(trapl true) -> trapl with t (trapl true)
    # puts spaces around braces so they become separate tokens
    for o, s in {'(': ' ( ', ')': ' ) ', '@': 'trapl with '}.items():
        code = code.replace(o, s) # also handles a tricky-to-use shortcut
    return code

def trapl_eval(code, syntax=None, unbox=True): # evaluate a string
    r = _trapl_eval(parse((syntax or syntax_plain)(code)))
    return r._val_ if unbox else r

# Now that we have a bare language to play with,
# let's use it to extend its standard library!
TRAPL = trapl_eval('trapl ext trapl true ( trapl false not )', unbox=False)
treval = lambda code: trapl_eval(code, syntax=syntax_rich, unbox=False) # short
TRAPL = treval("trapl.ext trapl 'ign' {x|x}") # trapl.ign x -> x
TRAPL = treval("""
mint = (trapl.atch trapl.int 'dec' {x|x neg inc neg})
mint = (trapl.atch mint 'sub' {x y|y neg add x})
mint = (trapl.atch mint 'lt' {x y|x ge y not})
mint = (trapl.atch mint 'le' {x y|y ge x})
mint = (trapl.atch mint 'gt' {x y|y ge x not})
mod_ = {a d|
 trapl.eval (trapl.if (a ge d) (
   trapl.code mod_ (a sub d) d
 ) (
   trapl.code a
 ))
}
mint = (trapl.atch mint 'mod' {a d|mod_ a d})
trapl.ext trapl 'int' mint
""") # NOTE: mod only works for small positive numbers

if __name__ == '__main__':
    import sys # evaluates a list of files or stdin contents
    for f in [file(fname) for fname in sys.argv[1:]] or [sys.stdin]:
        print trapl_eval(f.read(), syntax=syntax_rich)
<commit_msg>make mod use trapl.int.mod, not mod_ from a closure<commit_after>#!/usr/bin/python
# Toy Reduced Awful Programming Language
# Copyright (c) 2014 Alexander Sosedkin <monk@unboiled.info>, see LICENSE file

# This program is an interpreter for a toy language with simple syntax rules:
# (try reading the examples in test.py, most are pretty self-explanatory!)
# 1) Braces define priority
# 2) Applications chains: 'obj1 obj2 obj3' means obj3 gets applied to
#    the result of obj2 to obj1 application and so on
# Application usually means a dictionary lookup in first object,
# but there also are other important effects of application (see trapl_apply):
# 1) if first is a 'callable' it calls a wrapped python functions of 1 argument
# 2) if second is a 'method' it gets applied to the first object
# Additional syntax is supported with source-to-source transformations.
# The interpreter reuses python stack, closures... staying relatively compact.
# For even simpler, but less capable versions, please refer to the previous
# versions in the VCS (first version was very simple and < 87 lines long!).

class TRAPLError(RuntimeError): pass # a base exception for interpreter errors

class mydict(dict): # a lousily implemented immutable dictionary
    def __getattr__(self, k):    return self[k]
    def __setitem__(self, k, v): raise TRAPLError('__setitem__ disabled')
    def __setattr__(self, k, v): raise TRAPLError('__setattr__ disabled')
    def clear(self, *a, **kwar): raise TRAPLError('clear disabled')
    def update(self, *a, **kwa): raise TRAPLError('update disabled')
    def setdefault(s, *a, **kw): raise TRAPLError('setdefault disabled')
    def __call__(self, *a, **k): # this is how you extend mydict:
        n = self.copy()          # mydict({'a': 'b'})({'c': 'd'})(x='y')
        n.update(dict(*a, **k))  # results in a new mydict equivalent to
        return mydict(n)         # mydict({'a': 'b', 'c': 'd', 'x': 'y'})

OBJ = mydict() # create an immutable object, the base for all future objects
# methods and callables are objects with special fields (see trapl_apply)
METH = lambda body, **kwa: OBJ(_meth_=body, **kwa) # method and callables are
CALL = lambda body, **kwa: OBJ(_call_=body, **kwa) # created like this
# this is how you create a method of obj a that takes an additional parameter:
EQ = METH(lambda a: CALL(lambda b: # this means the method returns a callable
    TRAPL['true'] if a._val_ == b._val_ else TRAPL['false'] # which returns this
))
TRAPL = OBJ({ # let's define the standard library, accessible later as 'trapl'
    'obj': OBJ, # provide a means of accessing the original, blank object
    'false': OBJ({ # an instance of a boolean value false
        '_val_': False, # put a python-object value in _val_
        'not': METH(lambda b: b(_val_=(not b._val_))), # and several methods
        # understanding how they work remains as an exercise
        'or': METH(lambda b: CALL(lambda arg:
            arg if not b._val_ else b # this way 'false or 2' is 2
        )),
        'and': METH(lambda b: CALL(lambda arg:
            arg if b._val_ else b # this way 'true and 2' is 2
        )),
        'eq': EQ, # this is the example universal method from the above
    }),
    'int': OBJ({
        '_val_': 0,
        'new': METH(lambda i: CALL(lambda s: i(_val_=int(s._val_)))),
        'neg': METH(lambda i: i(_val_=-i._val_)),
        'inc': METH(lambda i: i(_val_=(i._val_ + 1))),
        'add': METH(lambda a: CALL(lambda b: a(_val_=a._val_ + b._val_))),
        'str': METH(lambda a: TRAPL['str'](_val_=str(a._val_))),
        'ge': METH(lambda a: CALL(lambda b:
            TRAPL['true' if a._val_ >= b._val_ else 'false']
        )),
        'eq': EQ,
    }),
    'str': OBJ({
        '_val_': '',
        'new': METH(lambda s: CALL(lambda n: s(_val_=n._val_))),
        'cat': METH(lambda s: CALL(lambda n: s(_val_=(s._val_ + n._val_)))),
        'len': METH(lambda s: TRAPL['int'](_val_=len(s._val_))),
        'rev': METH(lambda s: s(_val_=s._val_[::-1])),
        'dec': METH(lambda s: CALL(lambda e: s(_val_=decode_str(e._val_)))),
        'eq': EQ,
    }),
    'ext': CALL(lambda o: CALL(lambda n: CALL(lambda w: o({n._val_: w})))),
    'code': OBJ(_magic_='code'), # objects with _magic_ are special
    'eval': CALL(lambda code: OBJ(_magic_='eval', _magic_code_=code)),
    'drop': CALL(lambda name: OBJ(_magic_='drop', _magic_name_=name._val_)),
    'with': CALL(lambda name: CALL(lambda val:
        OBJ(_magic_='with', _magic_name_=name._val_, _magic_value_=val)
    )), # injects {name._val_: val} in current context, see trapl_eval
    'get': CALL(lambda name: OBJ(_magic_='get', _magic_name_=name._val_)),
    'func': CALL(lambda arg_name: CALL(lambda code:
        OBJ(_magic_='func', _magic_code_=code, _magic_arg_name_=arg_name),
    )), # make a callable will evaluate code with {arg_name._val_: arg_val}
    'atch': CALL(lambda o: CALL(lambda name: CALL(lambda call:
            o({name._val_: METH(call._call_)})
    ))), # Attaches a callable call to an object o as amethod with name name
    'dtch': CALL(lambda o: CALL(lambda name:
        CALL(o[name._val_]._meth_)
    )), # Returns a method detached from an object and turned into a callable
    'if': CALL(lambda cond:
        CALL(lambda val_true: CALL(lambda unused_val_false: val_true))
        if cond._val_ == True else
        CALL(lambda unused_val_true: CALL(lambda val_false: val_false))
    ),
})

def parse(tokens): # turns 'a ( b ( c ) d )' into ['a', ['b', ['c'], 'd']]
    if isinstance(tokens, str): tokens = tokens.split()
    tree = []
    while tokens:
        t = tokens.pop(0)
        if t == '(': tree.append(parse(tokens))
        elif t == ')': break
        else: tree.append(t)
    return tree

# Does the reverse job
flatten = lambda t: \
    ' '.join('( ' + flatten(a) + ' )' if isinstance(a, list) else a for a in t)

def trapl_apply(to, what): # Handles how one object gets applied to another
    if '_call_' in to: return to._call_(what) # appying to a callable calls it
    if not '_val_' in what: raise TRAPLError('No _val_ in %s' % what)
    if what._val_ in to: # if it's a string mathcing one of the object's keys
        x = to[what._val_] # retrieve it from the object
    # TODO LATER: Implement this functionality in trapl, remove from here
    elif '_get_' in to: # otherwise call a special catchall callable if present
        x = trapl_apply(trapl_apply(to, TRAPL['str'](_val_='_get_')), what)
    else:
        raise TRAPLError('No %s in %s' % (what._val_, to.keys()))
    if '_meth_' in x: return x._meth_(to) # if the result is a method, apply it
    # Note: to get a method from the object without appying it, see dtach
    return x

def _trapl_eval(tree, context=None, default_object=OBJ): # evaluate a tree
    context = {'trapl': TRAPL} if context is None else context.copy()
    curr = None
    while tree:
        next, tree = tree[0], tree[1:]
        if isinstance(next, list): next = _trapl_eval(next, context)
        if isinstance(next, str): # Autocast unknown literals to strings
            next = context.get(next, TRAPL['str'](_val_=next))
        curr = trapl_apply(curr, next) if not curr is None else next
        while '_magic_' in curr: # Handle special messages to the interpreter
            if curr._magic_ == 'with': # inject a value in current context
                context[curr._magic_name_] = curr._magic_value_
                curr = TRAPL['ign']
            elif curr._magic_ == 'drop': # inject a value in current context
                del context[curr._magic_name_]
                curr = TRAPL['ign']
            elif curr._magic_ == 'get': # get a value from current context
                curr = context[curr._magic_name_] # see '4+' example in test.py
            elif curr._magic_ == 'eval': # evaluate a string
                utree = parse(curr._magic_code_._val_)
                ctx = context.copy()
                if '_magic_definition_context_' in curr:
                    ctx.update(curr._magic_definition_context_)
                if '_magic_arg_val_' in curr:
                    ctx[curr._magic_arg_name_._val_] = curr._magic_arg_val_
                curr = _trapl_eval(utree, ctx)
            elif curr._magic_ == 'func': # create a function (closure)
                ctx, arg_name = context.copy(), curr._magic_arg_name_
                code = curr._magic_code_
                curr = CALL(lambda arg_val:
                    OBJ(_magic_='eval',
                        _magic_code_=code,
                        _magic_arg_name_=arg_name,
                        _magic_arg_val_=arg_val,
                        _magic_definition_context_=context,
                ))
            elif curr._magic_ == 'code': # create a string from remaining code
                # fall out of current brace if empty (allows trapl.code)
                if not tree: return curr
                curr = TRAPL['str'](_val_=flatten(tree))
                tree = []
    return curr or default_object

syntax_plain = lambda code: code # source code transformations may be used
# Let's define some code transformations to make syntax more pleasant
import re, base64
encode_str = lambda s: 'ENC' + base64.b32encode(s).replace('=', '0')
decode_str = lambda s: base64.b32decode(s[3:].replace('0', '='))
include_str = lambda s: ' ( trapl str dec %s ) ' % encode_str(s)
quotes = lambda code: re.sub(r"\B'([^'\\]*(?:\\.[^'\\]*)*)'\B", lambda m:
    include_str(m.group(1).replace('\\\'', '\'').replace('\\\\', '\\')), code)
dots = lambda code: re.sub(r"(\w*(?:\.(?:\w*))+)", lambda m:
    ' ( ' + m.group(1).split('.')[0] + ' ' +
    ' '.join(include_str(s) for s in m.group(1).split('.')[1:]) + ' ) ', code)
assign = lambda code: re.sub(r"(\w*)(?:\s*)=", lambda m:
    'trapl with %s ' % include_str(m.group(1)), code)

curly_func = lambda code: flatten(_curly_func(parse(
    code.replace('{', ' ( { ').replace('}', ' } ) ').replace('|', ' | ')
)))
funcize = lambda argname, tree: \
    ['trapl', 'func', parse(include_str(argname)), ['trapl', 'code'] + tree]
def _curly_func(tree):
    if isinstance(tree, str): return tree
    if len(tree) > 3:
        if tree[0] == '{' and '|' in tree[1:-1] and tree[-1] == '}':
            i = tree[1:-1].index('|') + 1
            args, tree = tree[1:i], tree[i+1:-1]
            for a in args[::-1]:
                tree = funcize(a, tree)
    return [_curly_func(t) for t in tree]

autoint = lambda code: re.sub(r"\b(\d+)\b", lambda m:
    ' ( trapl int new ' + include_str(m.group(1)) + ' ) ', code)

def syntax_rich(code): # apply lots of source-to-source transformations
    # convert 'a' into ( trapl string dec ENCSMTH ) to protect it
    code = quotes(code) # from futher damage, respects escaping
    # convert {a b|b a} -> (trapl func a ( trapl code ( trapl func b (
    code = code.replace('(', ' ( ').replace(')', ' ) ') # HACK: deduplicate
    code = curly_func(code) # trapl code b a ) ) ) )
    code = autoint(code) # convert 3 -> ( trapl int new 3 )
    code = dots(code) # convert t = trapl.true -> t = (trapl true)
    code = assign(code) # convert t=(trapl true) -> trapl with t (trapl true)
    # puts spaces around braces so they become separate tokens
    for o, s in {'(': ' ( ', ')': ' ) ', '@': 'trapl with '}.items():
        code = code.replace(o, s) # also handles a tricky-to-use shortcut
    return code

def trapl_eval(code, syntax=None, unbox=True): # evaluate a string
    r = _trapl_eval(parse((syntax or syntax_plain)(code)))
    return r._val_ if unbox else r

# Now that we have a bare language to play with,
# let's use it to extend its standard library!
TRAPL = trapl_eval('trapl ext trapl true ( trapl false not )', unbox=False)
treval = lambda code: trapl_eval(code, syntax=syntax_rich, unbox=False) # short
TRAPL = treval("trapl.ext trapl 'ign' {x|x}") # trapl.ign x -> x
TRAPL = treval("""
mint = (trapl.atch trapl.int 'dec' {x|x neg inc neg})
mint = (trapl.atch mint 'sub' {x y|y neg add x})
mint = (trapl.atch mint 'lt' {x y|x ge y not})
mint = (trapl.atch mint 'le' {x y|y ge x})
mint = (trapl.atch mint 'gt' {x y|y ge x not})
mod = {a d|
 trapl.eval (trapl.if (a ge d) (
   trapl.code (a sub d) mod d
 ) (
   trapl.code a
 ))
}
mint = (trapl.atch mint 'mod' mod)
trapl.ext trapl 'int' mint
""") # NOTE: mod only works for small positive numbers

if __name__ == '__main__':
    import sys # evaluates a list of files or stdin contents
    for f in [file(fname) for fname in sys.argv[1:]] or [sys.stdin]:
        print trapl_eval(f.read(), syntax=syntax_rich)
<|endoftext|>


<commit_before># coding=utf-8

from __future__ import print_function
from __future__ import unicode_literals

import inspect
import json
import os
import socket
import stat
import string
import sys

from builtins import int
from builtins import object
from builtins import open
from builtins import range

from .error import (InterfaceNotFound, InvalidParameter, MethodNotImplemented, VarlinkEncoder, VarlinkError,
                    ConnectionError)
from .scanner import Interface

try:
    from socketserver import (StreamRequestHandler, BaseServer, ThreadingMixIn, ForkingMixIn)
except ImportError: # Python2
    from SocketServer import (StreamRequestHandler, BaseServer, ThreadingMixIn, ForkingMixIn)

from types import GeneratorType

class Service(object):
    """Varlink service server handler

    To use the Service, a global object is instantiated:
    >>> service = Service(
    >>>    vendor='Red Hat',
    >>>    product='Manage System Accounts',
    >>>    version='1',
    >>>    interface_dir=os.path.dirname(__file__)
    >>> )

    For the class implementing the methods of a specific varlink interface
    a decorator is used:
    @service.interface('com.redhat.system.accounts')
    class Accounts:
    […]

    The varlink file corresponding to this interface is loaded from the 'interface_dir'
    specified in the constructor of the Service. It has to end in '.varlink'.

    Split the incoming stream for every null byte and feed it to the service.handle()
    function. Write any message returned from this generator function to the output stream.
    for outgoing_message in service.handle(incoming_message):
        connection.write(outgoing_message)

    or see, how the RequestHandler handles the Service object.

    Note: varlink only handles one method call at a time on one connection.

    """

    def __init__(self, vendor='', product='', version='', url='', interface_dir='.', namespaced=False):
        """Initialize the service with the data org.varlink.service.GetInfo() returns

        Arguments:
        interface_dir -- the directory with the *.varlink files for the interfaces
        """
        self.vendor = vendor
        self.product = product
        self.version = version
        self.url = url
        self.interface_dir = interface_dir
        self._namespaced = namespaced

        self.interfaces = {}
        self.interfaces_handlers = {}
        directory = os.path.dirname(__file__)
        self._add_interface(os.path.join(directory, 'org.varlink.service.varlink'), self)

    def GetInfo(self):
        """The standardized org.varlink.service.GetInfo() varlink method."""
        return {
            'vendor': self.vendor,
            'product': self.product,
            'version': self.version,
            'url': self.url,
            'interfaces': list(self.interfaces.keys())
        }

    def GetInterfaceDescription(self, interface):
        """The standardized org.varlink.service.GetInterfaceDescription() varlink method."""
        try:
            i = self.interfaces[interface]
        except KeyError:
            raise InterfaceNotFound(interface)

        return {'description': i.description}

    def _handle(self, message, raw_message, _server=None, _request=None):
        try:
            interface_name, _, method_name = message.get('method', '').rpartition('.')
            if not interface_name or not method_name:
                raise InterfaceNotFound(interface_name)

            interface = self.interfaces.get(interface_name)
            if not interface:
                raise InterfaceNotFound(interface_name)

            method = interface.get_method(method_name)

            parameters = message.get('parameters', {})
            handler = self.interfaces_handlers[interface.name]

            for name in parameters:
                if name not in method.in_type.fields:
                    raise InvalidParameter(name)

            for name in method.in_type.fields:
                if name not in parameters:
                    parameters[name] = None

            parameters = interface.filter_params("server.call", method.in_type, self._namespaced, parameters, None)

            func = getattr(handler, method_name, None)

            if not func or not callable(func):
                raise MethodNotImplemented(method_name)

            kwargs = {}

            if hasattr(inspect, "signature"):
                sig = inspect.signature(func)
                arg_names = [ (sig.parameters[k].kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) and k or None) for k in sig.parameters.keys()]
            else:
                from itertools import izip
                spec = inspect.getargspec(func)
                matched_args = [reversed(x) for x in [spec.args, spec.defaults or []]]
                arg_names = dict(izip(*matched_args))

            if message.get('more', False) or message.get('oneway', False) or message.get('upgrade', False):
                if message.get('more', False) and '_more' in arg_names:
                    kwargs["_more"] = True

                if message.get('oneway', False) and '_oneway' in arg_names:
                    kwargs["_oneway"] = True

                if message.get('upgrade', False) and '_upgrade' in arg_names:
                    kwargs["_upgrade"] = True

            if '_raw' in arg_names:
                kwargs["_raw"] = raw_message
            if '_message' in arg_names:
                kwargs["_message"] = message
            if '_interface' in arg_names:
                kwargs["_interface"] = interface
            if '_method' in arg_names:
                kwargs["_method"] = method
            if '_server' in arg_names:
                kwargs["_server"] = _server
            if '_request' in arg_names:
                kwargs["_request"] = _request

            if self._namespaced:
                out = func(*(getattr(parameters, k) for k in method.in_type.fields.keys()), **kwargs)
            else:
                out = func(*(parameters[k] for k in method.in_type.fields.keys()), **kwargs)

            if isinstance(out, GeneratorType):
                try:
                    for o in out:
                        if isinstance(o, Exception):
                            raise o

                        if kwargs.get("_oneway", False):
                            continue

                        cont = True
                        if '_continues' in o:
                            cont = o['_continues']
                            del o['_continues']
                            yield {'continues': bool(cont), 'parameters': interface.filter_params("server.reply", method.out_type,
                                                                                self._namespaced, o,
                                                                                None) or {}}
                        else:
                            yield {'parameters': interface.filter_params("server.reply", method.out_type,
                                                                         self._namespaced, o,
                                                                         None) or {}}

                        if not cont:
                            return
                except ConnectionError as e:
                    try:
                        out.throw(e)
                    except StopIteration:
                        pass
            else:
                if message.get('oneway', False):
                    yield None
                else:
                    yield {'parameters': out or {}}

        except VarlinkError as error:
            yield error

    def handle(self, message, _server=None, _request=None):
        """This generator function handles any incoming message. Write any returned bytes to the output stream.

        for outgoing_message in service.handle(incoming_message):
            connection.write(outgoing_message)
        """
        if not message:
            return

        if message[-1] == 0:
            message = message[:-1]

        handle = self._handle(json.loads(message), message, _server, _request)
        for out in handle:
            if out == None:
                return
            try:
                yield json.dumps(out, cls=VarlinkEncoder).encode('utf-8')
            except ConnectionError as e:
                try:
                    handle.throw(e)
                except StopIteration:
                    pass

    def _add_interface(self, filename, handler):
        if not os.path.isabs(filename):
            filename = os.path.join(self.interface_dir, filename + '.varlink')

        with open(filename) as f:
            interface = Interface(f.read())
            self.interfaces[interface.name] = interface
            self.interfaces_handlers[interface.name] = handler

    def interface(self, filename):

        def decorator(interface_class):
            self._add_interface(filename, interface_class())
            return interface_class

        return decorator

def get_listen_fd():
    if "LISTEN_FDS" not in os.environ:
        return None
    if "LISTEN_PID" not in os.environ:
        return None
    try:
        if int(os.environ["LISTEN_PID"]) != os.getpid():
            return None
    except:
        return None

    try:
        fds = int(os.environ["LISTEN_FDS"])
    except:
        return None

    if fds < 1:
        return None

    if fds == 1:
        try:
            if stat.S_ISSOCK(os.fstat(3).st_mode):
                return 3
            else:
                return None
        except OSError:
            return None

    fields = string.split(os.environ["LISTEN_FDNAMES"], ":")

    if len(fields) != fds:
        return None

    for i in range(len(fields)):
        if fields[i] == "varlink":
            try:
                if stat.S_ISSOCK(os.fstat(i + 3).st_mode):
                    return i + 3
                else:
                    return None
            except OSError:
                return None

class RequestHandler(StreamRequestHandler):
    """Varlink request handler

    To use as an argument for the VarlinkServer constructor.
    Instantiate your own class and set the class variable service to your global varlink.Service object.
    """
    service = None

    def handle(self):
        message = b''

        self.request.setblocking(True)
        while not self.rfile.closed:
            c = self.rfile.read(1)

            if c == b'':
                break

            if c != b'\0':
                message += c
                continue

            for reply in self.service.handle(message, _server=self.server, _request=self.request):
                if reply != None:
                    self.wfile.write(reply + b'\0')

            message = b''

class Server(BaseServer):
    """Server

    The same as the standard socketserver.TCPServer, to initialize with a subclass of RequestHandler.

    >>> import varlink
    >>> import os
    >>>
    >>> service = varlink.Service(vendor='Example', product='Examples', version='1', url='http://example.com',
    >>>    interface_dir=os.path.dirname(__file__))
    >>>
    >>> class ServiceRequestHandler(varlink.RequestHandler):
    >>>    service = service
    >>>
    >>> @service.interface('com.example.service')
    >>> class Example:
    >>>    # com.example.service method implementation here …
    >>>    pass
    >>>
    >>> server = varlink.ThreadingServer(sys.argv[1][10:], ServiceRequestHandler)
    >>> server.serve_forever()
    """
    address_family = socket.AF_INET

    socket_type = socket.SOCK_STREAM

    request_queue_size = 5

    allow_reuse_address = True

    def __init__(self, server_address, RequestHandlerClass, bind_and_activate=True):
        self.remove_file = None
        self.mode = None
        self.listen_fd = get_listen_fd()

        if self.listen_fd:
            server_address = self.listen_fd
            self.address_family = socket.AF_UNIX
            self.socket = socket.fromfd(self.listen_fd, socket.AF_UNIX, socket.SOCK_STREAM)

        elif server_address.startswith("unix:"):
            self.address_family = socket.AF_UNIX
            address = server_address[5:]
            m = address.rfind(';mode=')
            if m != -1:
                self.mode = address[m + 6:]
                address = address[:m]

            if address[0] == '@':
                address = address.replace('@', '\0', 1)
                self.mode = None
            else:
                self.remove_file = address

            server_address = address
            self.socket = socket.socket(self.address_family, self.socket_type)
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        elif server_address.startswith("tcp:"):
            address = server_address[4:]
            p = address.rfind(':')
            if p != -1:
                port = int(address[p + 1:])
                address = address[:p]
            else:
                raise ConnectionError("Invalid address 'tcp:%s'" % address)
            address = address.replace('[', '')
            address = address.replace(']', '')

            try:
                res = socket.getaddrinfo(address, port, proto=socket.IPPROTO_TCP, flags=socket.AI_NUMERICHOST)
            except TypeError:
                res = socket.getaddrinfo(address, port, self.address_family, self.socket_type, socket.IPPROTO_TCP, socket.AI_NUMERICHOST)

            af, socktype, proto, canonname, sa = res[0]
            self.address_family = af
            self.socket_type = socktype
            self.socket = socket.socket(self.address_family, self.socket_type)
            server_address = sa[0:2]

        else:
            raise ConnectionError("Invalid address '%s'" % server_address)

        BaseServer.__init__(self, server_address, RequestHandlerClass)

        if bind_and_activate:
            try:
                self.server_bind()
                self.server_activate()
            except:
                self.server_close()
                raise

    def server_bind(self):
        """Called by constructor to bind the socket.

        May be overridden.

        """
        if self.allow_reuse_address:
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.socket.setblocking(True)

        if not self.listen_fd:
            self.socket.bind(self.server_address)

        self.server_address = self.socket.getsockname()

        if self.server_address[0] == 0:
            self.server_address = '@' + self.server_address[1:].decode('utf-8')
        elif self.mode:
            os.chmod(self.server_address, mode=int(self.mode, 8))

    def server_activate(self):
        """Called by constructor to activate the server.

        May be overridden.

        """
        self.socket.listen(self.request_queue_size)

    def server_close(self):
        """Called to clean-up the server.

        May be overridden.

        """
        if self.remove_file:
            try:
                os.remove(self.remove_file)
            except:
                pass
        self.socket.close()

    def fileno(self):
        """Return socket file number.

        Interface required by selector.

        """
        return self.socket.fileno()

    def get_request(self):
        """Get the request and client address from the socket.

        May be overridden.

        """
        return self.socket.accept()

    def shutdown_request(self, request):
        """Called to shutdown and close an individual request."""
        try:
            # explicitly shutdown.  socket.close() merely releases
            # the socket and waits for GC to perform the actual close.
            request.shutdown(socket.SHUT_WR)
        except:
            pass  # some platforms may raise ENOTCONN here
        self.close_request(request)

    def close_request(self, request):
        """Called to clean up an individual request."""
        request.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.server_close()

class ThreadingServer(ThreadingMixIn, Server): pass

class ForkingServer(ForkingMixIn, Server): pass
<commit_msg>on server shutdown_request, also shutdown read<commit_after># coding=utf-8

from __future__ import print_function
from __future__ import unicode_literals

import inspect
import json
import os
import socket
import stat
import string
import sys

from builtins import int
from builtins import object
from builtins import open
from builtins import range

from .error import (InterfaceNotFound, InvalidParameter, MethodNotImplemented, VarlinkEncoder, VarlinkError,
                    ConnectionError)
from .scanner import Interface

try:
    from socketserver import (StreamRequestHandler, BaseServer, ThreadingMixIn, ForkingMixIn)
except ImportError: # Python2
    from SocketServer import (StreamRequestHandler, BaseServer, ThreadingMixIn, ForkingMixIn)

from types import GeneratorType

class Service(object):
    """Varlink service server handler

    To use the Service, a global object is instantiated:
    >>> service = Service(
    >>>    vendor='Red Hat',
    >>>    product='Manage System Accounts',
    >>>    version='1',
    >>>    interface_dir=os.path.dirname(__file__)
    >>> )

    For the class implementing the methods of a specific varlink interface
    a decorator is used:
    @service.interface('com.redhat.system.accounts')
    class Accounts:
    […]

    The varlink file corresponding to this interface is loaded from the 'interface_dir'
    specified in the constructor of the Service. It has to end in '.varlink'.

    Split the incoming stream for every null byte and feed it to the service.handle()
    function. Write any message returned from this generator function to the output stream.
    for outgoing_message in service.handle(incoming_message):
        connection.write(outgoing_message)

    or see, how the RequestHandler handles the Service object.

    Note: varlink only handles one method call at a time on one connection.

    """

    def __init__(self, vendor='', product='', version='', url='', interface_dir='.', namespaced=False):
        """Initialize the service with the data org.varlink.service.GetInfo() returns

        Arguments:
        interface_dir -- the directory with the *.varlink files for the interfaces
        """
        self.vendor = vendor
        self.product = product
        self.version = version
        self.url = url
        self.interface_dir = interface_dir
        self._namespaced = namespaced

        self.interfaces = {}
        self.interfaces_handlers = {}
        directory = os.path.dirname(__file__)
        self._add_interface(os.path.join(directory, 'org.varlink.service.varlink'), self)

    def GetInfo(self):
        """The standardized org.varlink.service.GetInfo() varlink method."""
        return {
            'vendor': self.vendor,
            'product': self.product,
            'version': self.version,
            'url': self.url,
            'interfaces': list(self.interfaces.keys())
        }

    def GetInterfaceDescription(self, interface):
        """The standardized org.varlink.service.GetInterfaceDescription() varlink method."""
        try:
            i = self.interfaces[interface]
        except KeyError:
            raise InterfaceNotFound(interface)

        return {'description': i.description}

    def _handle(self, message, raw_message, _server=None, _request=None):
        try:
            interface_name, _, method_name = message.get('method', '').rpartition('.')
            if not interface_name or not method_name:
                raise InterfaceNotFound(interface_name)

            interface = self.interfaces.get(interface_name)
            if not interface:
                raise InterfaceNotFound(interface_name)

            method = interface.get_method(method_name)

            parameters = message.get('parameters', {})
            handler = self.interfaces_handlers[interface.name]

            for name in parameters:
                if name not in method.in_type.fields:
                    raise InvalidParameter(name)

            for name in method.in_type.fields:
                if name not in parameters:
                    parameters[name] = None

            parameters = interface.filter_params("server.call", method.in_type, self._namespaced, parameters, None)

            func = getattr(handler, method_name, None)

            if not func or not callable(func):
                raise MethodNotImplemented(method_name)

            kwargs = {}

            if hasattr(inspect, "signature"):
                sig = inspect.signature(func)
                arg_names = [ (sig.parameters[k].kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) and k or None) for k in sig.parameters.keys()]
            else:
                from itertools import izip
                spec = inspect.getargspec(func)
                matched_args = [reversed(x) for x in [spec.args, spec.defaults or []]]
                arg_names = dict(izip(*matched_args))

            if message.get('more', False) or message.get('oneway', False) or message.get('upgrade', False):
                if message.get('more', False) and '_more' in arg_names:
                    kwargs["_more"] = True

                if message.get('oneway', False) and '_oneway' in arg_names:
                    kwargs["_oneway"] = True

                if message.get('upgrade', False) and '_upgrade' in arg_names:
                    kwargs["_upgrade"] = True

            if '_raw' in arg_names:
                kwargs["_raw"] = raw_message
            if '_message' in arg_names:
                kwargs["_message"] = message
            if '_interface' in arg_names:
                kwargs["_interface"] = interface
            if '_method' in arg_names:
                kwargs["_method"] = method
            if '_server' in arg_names:
                kwargs["_server"] = _server
            if '_request' in arg_names:
                kwargs["_request"] = _request

            if self._namespaced:
                out = func(*(getattr(parameters, k) for k in method.in_type.fields.keys()), **kwargs)
            else:
                out = func(*(parameters[k] for k in method.in_type.fields.keys()), **kwargs)

            if isinstance(out, GeneratorType):
                try:
                    for o in out:
                        if isinstance(o, Exception):
                            raise o

                        if kwargs.get("_oneway", False):
                            continue

                        cont = True
                        if '_continues' in o:
                            cont = o['_continues']
                            del o['_continues']
                            yield {'continues': bool(cont), 'parameters': interface.filter_params("server.reply", method.out_type,
                                                                                self._namespaced, o,
                                                                                None) or {}}
                        else:
                            yield {'parameters': interface.filter_params("server.reply", method.out_type,
                                                                         self._namespaced, o,
                                                                         None) or {}}

                        if not cont:
                            return
                except ConnectionError as e:
                    try:
                        out.throw(e)
                    except StopIteration:
                        pass
            else:
                if message.get('oneway', False):
                    yield None
                else:
                    yield {'parameters': out or {}}

        except VarlinkError as error:
            yield error

    def handle(self, message, _server=None, _request=None):
        """This generator function handles any incoming message. Write any returned bytes to the output stream.

        for outgoing_message in service.handle(incoming_message):
            connection.write(outgoing_message)
        """
        if not message:
            return

        if message[-1] == 0:
            message = message[:-1]

        handle = self._handle(json.loads(message), message, _server, _request)
        for out in handle:
            if out == None:
                return
            try:
                yield json.dumps(out, cls=VarlinkEncoder).encode('utf-8')
            except ConnectionError as e:
                try:
                    handle.throw(e)
                except StopIteration:
                    pass

    def _add_interface(self, filename, handler):
        if not os.path.isabs(filename):
            filename = os.path.join(self.interface_dir, filename + '.varlink')

        with open(filename) as f:
            interface = Interface(f.read())
            self.interfaces[interface.name] = interface
            self.interfaces_handlers[interface.name] = handler

    def interface(self, filename):

        def decorator(interface_class):
            self._add_interface(filename, interface_class())
            return interface_class

        return decorator

def get_listen_fd():
    if "LISTEN_FDS" not in os.environ:
        return None
    if "LISTEN_PID" not in os.environ:
        return None
    try:
        if int(os.environ["LISTEN_PID"]) != os.getpid():
            return None
    except:
        return None

    try:
        fds = int(os.environ["LISTEN_FDS"])
    except:
        return None

    if fds < 1:
        return None

    if fds == 1:
        try:
            if stat.S_ISSOCK(os.fstat(3).st_mode):
                return 3
            else:
                return None
        except OSError:
            return None

    fields = string.split(os.environ["LISTEN_FDNAMES"], ":")

    if len(fields) != fds:
        return None

    for i in range(len(fields)):
        if fields[i] == "varlink":
            try:
                if stat.S_ISSOCK(os.fstat(i + 3).st_mode):
                    return i + 3
                else:
                    return None
            except OSError:
                return None

class RequestHandler(StreamRequestHandler):
    """Varlink request handler

    To use as an argument for the VarlinkServer constructor.
    Instantiate your own class and set the class variable service to your global varlink.Service object.
    """
    service = None

    def handle(self):
        message = b''

        self.request.setblocking(True)
        while not self.rfile.closed:
            c = self.rfile.read(1)

            if c == b'':
                break

            if c != b'\0':
                message += c
                continue

            for reply in self.service.handle(message, _server=self.server, _request=self.request):
                if reply != None:
                    self.wfile.write(reply + b'\0')

            message = b''

class Server(BaseServer):
    """Server

    The same as the standard socketserver.TCPServer, to initialize with a subclass of RequestHandler.

    >>> import varlink
    >>> import os
    >>>
    >>> service = varlink.Service(vendor='Example', product='Examples', version='1', url='http://example.com',
    >>>    interface_dir=os.path.dirname(__file__))
    >>>
    >>> class ServiceRequestHandler(varlink.RequestHandler):
    >>>    service = service
    >>>
    >>> @service.interface('com.example.service')
    >>> class Example:
    >>>    # com.example.service method implementation here …
    >>>    pass
    >>>
    >>> server = varlink.ThreadingServer(sys.argv[1][10:], ServiceRequestHandler)
    >>> server.serve_forever()
    """
    address_family = socket.AF_INET

    socket_type = socket.SOCK_STREAM

    request_queue_size = 5

    allow_reuse_address = True

    def __init__(self, server_address, RequestHandlerClass, bind_and_activate=True):
        self.remove_file = None
        self.mode = None
        self.listen_fd = get_listen_fd()

        if self.listen_fd:
            server_address = self.listen_fd
            self.address_family = socket.AF_UNIX
            self.socket = socket.fromfd(self.listen_fd, socket.AF_UNIX, socket.SOCK_STREAM)

        elif server_address.startswith("unix:"):
            self.address_family = socket.AF_UNIX
            address = server_address[5:]
            m = address.rfind(';mode=')
            if m != -1:
                self.mode = address[m + 6:]
                address = address[:m]

            if address[0] == '@':
                address = address.replace('@', '\0', 1)
                self.mode = None
            else:
                self.remove_file = address

            server_address = address
            self.socket = socket.socket(self.address_family, self.socket_type)
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        elif server_address.startswith("tcp:"):
            address = server_address[4:]
            p = address.rfind(':')
            if p != -1:
                port = int(address[p + 1:])
                address = address[:p]
            else:
                raise ConnectionError("Invalid address 'tcp:%s'" % address)
            address = address.replace('[', '')
            address = address.replace(']', '')

            try:
                res = socket.getaddrinfo(address, port, proto=socket.IPPROTO_TCP, flags=socket.AI_NUMERICHOST)
            except TypeError:
                res = socket.getaddrinfo(address, port, self.address_family, self.socket_type, socket.IPPROTO_TCP, socket.AI_NUMERICHOST)

            af, socktype, proto, canonname, sa = res[0]
            self.address_family = af
            self.socket_type = socktype
            self.socket = socket.socket(self.address_family, self.socket_type)
            server_address = sa[0:2]

        else:
            raise ConnectionError("Invalid address '%s'" % server_address)

        BaseServer.__init__(self, server_address, RequestHandlerClass)

        if bind_and_activate:
            try:
                self.server_bind()
                self.server_activate()
            except:
                self.server_close()
                raise

    def server_bind(self):
        """Called by constructor to bind the socket.

        May be overridden.

        """
        if self.allow_reuse_address:
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.socket.setblocking(True)

        if not self.listen_fd:
            self.socket.bind(self.server_address)

        self.server_address = self.socket.getsockname()

        if self.server_address[0] == 0:
            self.server_address = '@' + self.server_address[1:].decode('utf-8')
        elif self.mode:
            os.chmod(self.server_address, mode=int(self.mode, 8))

    def server_activate(self):
        """Called by constructor to activate the server.

        May be overridden.

        """
        self.socket.listen(self.request_queue_size)

    def server_close(self):
        """Called to clean-up the server.

        May be overridden.

        """
        if self.remove_file:
            try:
                os.remove(self.remove_file)
            except:
                pass
        self.socket.close()

    def fileno(self):
        """Return socket file number.

        Interface required by selector.

        """
        return self.socket.fileno()

    def get_request(self):
        """Get the request and client address from the socket.

        May be overridden.

        """
        return self.socket.accept()

    def shutdown_request(self, request):
        """Called to shutdown and close an individual request."""
        try:
            # explicitly shutdown.  socket.close() merely releases
            # the socket and waits for GC to perform the actual close.
            request.shutdown(socket.SHUT_RDWR)

        except:
            pass  # some platforms may raise ENOTCONN here
        self.close_request(request)

    def close_request(self, request):
        """Called to clean up an individual request."""
        request.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.server_close()

class ThreadingServer(ThreadingMixIn, Server): pass

class ForkingServer(ForkingMixIn, Server): pass
<|endoftext|>


<commit_before>import csv
from collections import OrderedDict
import json

from pymongo import MongoClient
from texttable import Texttable

from .logger import logger
from .bson_utils import get_dtype

# Export
__all__ = ['SchemaAnalyzer']

class SchemaAnalyzer(object):
    """SchemaAnalyzer is the main class for building
    a MongoDB schema analysis.

    The API is simple, instantiate a schema object:
    >>> schema = SchemaAnalyzer()
    >>> schema.info()

    You can export the schema to CSV:

    >>> schema.to_csv()

    The `to_csv()` method is just a wrapper around the
    csv writer method from the standard library, so you can
    pass in all the arguments you would pass to your csv writer.
    >>> schema.to_csv(delimiter=';')
    >>> schema.save()
    """
    def __init__(self, host="", db="", collection="", query={}, **kwargs):
        """
        TODO: solve working with new query
        The init metod provides some 'private' props, like `_len`, the
        length of the query results
        """
        if not host:
            self.cursor = MongoClient()
        else:
            self.cursor = MongoClient(host, **kwargs)
        self.db = db
        self.collection = collection
        self.query = query
        self.schema = {}
        self._len = 0

    def analyze(self):
        """
        TODO: easier collection selection
        TODO: make it poosible to analyze the db
        """
        with self.cursor as cursor:
            logger.info("Analyzing schema.")
            conn = cursor[self.db][self.collection]
            results = conn.find(self.query)
        # Analyze the results
        logger.debug("Analyzing {} in {}".format(self.collection, self.db))
        # TODO: Make a generator function
        for result in results:
            self._len += 1
            self._process_object(result)
        # Calculate the percentage for each field
        self._preprocesss_for_reproting()

    def _process_object(self, result):
        """
        """
        global curr_object
        curr_object = {}
        self._get_from_object(result)
        for key in curr_object.keys():
            if self.schema.get(key):
                self.schema[key]["sum"] += 1
            else:
                self.schema[key] = curr_object[key]

    def _get_from_object(self, result, path=[]):
        """
        The main method that creates the schema

        TODO: this is a first draft, a proof of concept
        TODO: the real thing must create a dict with
        TODO: dicts inside specifing the values and
        TODO: mitigatin the values + returning the types for js
        TODO: iteritems for python 2
        """
        for key, val in result.items():
            # logger.debug("Analyzing key {}".format(key))
            new_path = path + [key]
            if isinstance(val, dict):
                self._get_from_object(val, path=new_path)
            elif isinstance(val, list):
                self._get_from_list(val, path=new_path)
            else:
                full_path = '.'.join(new_path)
                data = curr_object.get(full_path)
                if not data:
                    curr_object[full_path] = {
                        'type': get_dtype(val),
                        'sum': 1
                    }

    def _get_from_list(self, results, path=[]):
        """
        Maps all elements of a list with the method `_get_from_object`.

        :param results: a list from a pymongo cursor
        :param path: a list with parent fields
        """
        for result in results:
            self._get_from_object(result, path=path)

    def _preprocesss_for_reproting(self):
        """Prepares the date for reporting to stdOut or JSON"""
        # Add percentage for field
        for value in self.schema.values():
            percentage = round(value["sum"]*100/self._len, 2)
            value["occurrence"] = str(percentage) + " %"
        # Prepare the data
        # first sorting by occurence
        # than changin occurence to %
        data = sorted(
            self.schema.items(),
            key=lambda x: x[1]["sum"],
            reverse=True
        )
        # Insert into OrderdDict
        # TODO exception for python 3.6 >
        self.schema = OrderedDict()
        for element in data:
            key, value = element
            self.schema[key] = value

    def _get_rows(self):
        """Returns array coresoinding to rows."""
        rows = []
        for key, value in self.schema.items():
            name = key
            type_ = value["type"]
            occurrence = value["occurrence"]
            rows.append([name, type_, occurrence])
        return rows

    def to_json(self):
        """JSON representation of the schema.
        TODO: the first line of to_json an __str__ duplicate code, refactor
        """
        if not self.schema:
            self.analyze()
        return json.dumps(self.schema)

    def to_csv(self, name="report.csv", **kwargs):
        """CSV representation of the schema."""
        if not self.schema:
            self.analyze()
        with open(name, "w") as f:
            csv_writer = csv.writer(f, **kwargs)
            csv_writer.writerow(["Field", "Data Type", "Occurrence"])
            rows = self._get_rows()
            for row in rows:
                csv_writer.writerow(row)

    def __str__(self, out="ascii"):
        """Printable representation of the schema."""
        if not self.schema:
            self.analyze()
        # Prepare the ASCII table
        table = Texttable()
        table.set_cols_align(['l', 'l', 'l'])
        table.set_cols_valign(['m', 'm', 'm'])
        table.set_cols_dtype(['t', 'i','a'])
        table.add_row(["Field", "Data Type", "Occurrence"])
        # Create the rows
        rows = self._get_rows()
        for row in rows:
            table.add_row(row)
        return table.draw() + '\n'
<commit_msg>New method to proces objects<commit_after>import csv
from collections import OrderedDict
import json

from pymongo import MongoClient
from texttable import Texttable

from .logger import logger
from .bson_utils import get_dtype

# Export
__all__ = ['SchemaAnalyzer']

class SchemaAnalyzer(object):
    """SchemaAnalyzer is the main class for building
    a MongoDB schema analysis.

    The API is simple, instantiate a schema object:
    >>> schema = SchemaAnalyzer()
    >>> schema.info()

    You can export the schema to CSV:

    >>> schema.to_csv()

    The `to_csv()` method is just a wrapper around the
    csv writer method from the standard library, so you can
    pass in all the arguments you would pass to your csv writer.
    >>> schema.to_csv(delimiter=';')
    >>> schema.save()
    """
    def __init__(self, host="", db="", collection="", query={}, **kwargs):
        """
        TODO: solve working with new query
        The init metod provides some 'private' props, like `_len`, the
        length of the query results
        """
        if not host:
            self.cursor = MongoClient()
        else:
            self.cursor = MongoClient(host, **kwargs)
        self.db = db
        self.collection = collection
        self.query = query
        self.schema = {}
        self._len = 0

    def analyze(self):
        """
        TODO: easier collection selection
        TODO: make it poosible to analyze the db
        """
        with self.cursor as cursor:
            logger.info("Analyzing schema.")
            conn = cursor[self.db][self.collection]
            results = conn.find(self.query)
        # Analyze the results
        logger.debug("Analyzing {} in {}".format(self.collection, self.db))
        # TODO: Make a generator function
        for result in results:
            self._len += 1
            self._process_object(result)
        # Calculate the percentage for each field
        self._preprocesss_for_reproting()

    """
    New and better method for single objects
     def get_keys(item):
    ...:     for k,v in item.items():
    ...:         if isinstance(v, dict):
    ...:             get_keys(v)
    ...:         elif isinstance(v, list):
    ...:             for e in v:
    ...:                 get_keys(e)
    ...:         else:
    ...:             if not i.get(k):
    ...:                 i[k] = 1
    """

    def _process_object(self, result):
        """
        """
        global curr_object
        curr_object = {}
        self._get_from_object(result)
        for key in curr_object.keys():
            if self.schema.get(key):
                self.schema[key]["sum"] += 1
            else:
                self.schema[key] = curr_object[key]

    def _get_from_object(self, result, path=[]):
        """
        The main method that creates the schema

        TODO: this is a first draft, a proof of concept
        TODO: the real thing must create a dict with
        TODO: dicts inside specifing the values and
        TODO: mitigatin the values + returning the types for js
        TODO: iteritems for python 2
        """
        for key, val in result.items():
            # logger.debug("Analyzing key {}".format(key))
            new_path = path + [key]
            if isinstance(val, dict):
                self._get_from_object(val, path=new_path)
            elif isinstance(val, list):
                self._get_from_list(val, path=new_path)
            else:
                full_path = '.'.join(new_path)
                data = curr_object.get(full_path)
                if not data:
                    curr_object[full_path] = {
                        'type': get_dtype(val),
                        'sum': 1
                    }

    def _get_from_list(self, results, path=[]):
        """
        Maps all elements of a list with the method `_get_from_object`.

        :param results: a list from a pymongo cursor
        :param path: a list with parent fields
        """
        for result in results:
            self._get_from_object(result, path=path)

    def _preprocesss_for_reproting(self):
        """Prepares the date for reporting to stdOut or JSON"""
        # Add percentage for field
        for value in self.schema.values():
            percentage = round(value["sum"]*100/self._len, 2)
            value["occurrence"] = str(percentage) + " %"
        # Prepare the data
        # first sorting by occurence
        # than changin occurence to %
        data = sorted(
            self.schema.items(),
            key=lambda x: x[1]["sum"],
            reverse=True
        )
        # Insert into OrderdDict
        # TODO exception for python 3.6 >
        self.schema = OrderedDict()
        for element in data:
            key, value = element
            self.schema[key] = value

    def _get_rows(self):
        """Returns array coresoinding to rows."""
        rows = []
        for key, value in self.schema.items():
            name = key
            type_ = value["type"]
            occurrence = value["occurrence"]
            rows.append([name, type_, occurrence])
        return rows

    def to_json(self):
        """JSON representation of the schema.
        TODO: the first line of to_json an __str__ duplicate code, refactor
        """
        if not self.schema:
            self.analyze()
        return json.dumps(self.schema)

    def to_csv(self, name="report.csv", **kwargs):
        """CSV representation of the schema."""
        if not self.schema:
            self.analyze()
        with open(name, "w") as f:
            csv_writer = csv.writer(f, **kwargs)
            csv_writer.writerow(["Field", "Data Type", "Occurrence"])
            rows = self._get_rows()
            for row in rows:
                csv_writer.writerow(row)

    def __str__(self, out="ascii"):
        """Printable representation of the schema."""
        if not self.schema:
            self.analyze()
        # Prepare the ASCII table
        table = Texttable()
        table.set_cols_align(['l', 'l', 'l'])
        table.set_cols_valign(['m', 'm', 'm'])
        table.set_cols_dtype(['t', 'i','a'])
        table.add_row(["Field", "Data Type", "Occurrence"])
        # Create the rows
        rows = self._get_rows()
        for row in rows:
            table.add_row(row)
        return table.draw() + '\n'
<|endoftext|>


<commit_before>#!/usr/bin/env python

from __future__ import print_function

import io
import os
import sys
import json
import time
import logging
import argparse
import requests

from requests_oauthlib import OAuth1Session

try:
    import configparser  # Python 3
except ImportError:
    import ConfigParser as configparser  # Python 2

if sys.version_info[:2] <= (2, 7):
    # Python 2
    get_input = raw_input
else:
    # Python 3
    get_input = input

def main():
    """
    The twarc command line.
    """
    parser = argparse.ArgumentParser("twarc")
    parser.add_argument("--search", dest="search",
                        help="search for tweets matching a query")
    parser.add_argument("--max_id", dest="max_id",
                        help="maximum tweet id to search for")
    parser.add_argument("--since_id", dest="since_id",
                        help="smallest id to search for")
    parser.add_argument("--stream", dest="stream",
                        help="stream tweets matching filter")
    parser.add_argument("--hydrate", dest="hydrate",
                        help="rehydrate tweets from a file of tweet ids")
    parser.add_argument("--log", dest="log",
                        default="twarc.log", help="log file")
    parser.add_argument("--consumer_key",
                        default=None, help="Twitter API consumer key")
    parser.add_argument("--consumer_secret",
                        default=None, help="Twitter API consumer secret")
    parser.add_argument("--access_token",
                        default=None, help="Twitter API access key")
    parser.add_argument("--access_token_secret",
                        default=None, help="Twitter API access token secret")
    parser.add_argument('-c', '--config',
                        default=default_config_filename(),
                        help="Config file containing Twitter keys and secrets")
    parser.add_argument('-p', '--profile', default='main',
                        help="Name of a profile in your configuration file")
    args = parser.parse_args()

    logging.basicConfig(
        filename=args.log,
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s"
    )

    consumer_key = args.consumer_key or os.environ.get('CONSUMER_KEY')
    consumer_secret = args.consumer_secret or os.environ.get('CONSUMER_SECRET')
    access_token = args.access_token or os.environ.get('ACCESS_TOKEN')
    access_token_secret = args.access_token_secret or os.environ.get("ACCESS_TOKEN_SECRET")

    if not (consumer_key and consumer_secret and
            access_token and access_token_secret):
        credentials = load_config(args.config, args.profile)
        if credentials:
            consumer_key = credentials['consumer_key']
            consumer_secret = credentials['consumer_secret']
            access_token = credentials['access_token']
            access_token_secret = credentials['access_token_secret']
        else:
            print("Please enter Twitter authentication credentials")
            consumer_key = get_input('consumer key: ')
            consumer_secret = get_input('consumer secret: ')
            access_token = get_input('access_token: ')
            access_token_secret = get_input('access token secret: ')
            save_keys(args.profile, consumer_key, consumer_secret,
                      access_token, access_token_secret)

    t = Twarc(consumer_key=consumer_key,
              consumer_secret=consumer_secret,
              access_token=access_token,
              access_token_secret=access_token_secret)

    if args.search:
        tweets = t.search(
            args.search,
            since_id=args.since_id,
            max_id=args.max_id
        )
    elif args.stream:
        tweets = t.stream(args.stream)
    elif args.hydrate:
        tweets = t.hydrate(open(args.hydrate, 'rU'))
    else:
        raise argparse.ArgumentTypeError(
            "must supply one of: --search --stream or --hydrate")

    # iterate through the tweets and write them to stdout
    for tweet in tweets:
        if "id_str" in tweet:
            logging.info("archived %s", tweet["id_str"])
            print(json.dumps(tweet))

def load_config(filename, profile):
    if not os.path.isfile(filename):
        return None
    config = configparser.ConfigParser()
    config.read(filename)
    data = {}
    for key in ['access_token', 'access_token_secret', 'consumer_key', 'consumer_secret']:
        try:
            data[key] = config.get(profile, key)
        except configparser.NoSectionError:
            sys.exit("no such profile %s in %s" % (profile, filename))
        except configparser.NoOptionError:
            sys.exit("missing %s from profile %s in %s" % (key, profile, filename))
    return data

def save_config(filename, profile,
                consumer_key, consumer_secret,
                access_token, access_token_secret):
    config = configparser.ConfigParser()
    config.add_section(profile)
    config.set(profile, 'consumer_key', consumer_key)
    config.set(profile, 'consumer_secret', consumer_secret)
    config.set(profile, 'access_token', access_token)
    config.set(profile, 'access_token_secret', access_token_secret)
    with open(filename, 'w') as config_file:
        config.write(config_file)

def default_config_filename():
    """
    Return the default filename for storing Twitter keys.
    """
    home = os.path.expanduser("~")
    return os.path.join(home, ".twarc")

def save_keys(profile, consumer_key, consumer_secret,
              access_token, access_token_secret):
    """
    Save keys to ~/.twarc
    """
    filename = default_config_filename()
    save_config(filename, profile,
                consumer_key, consumer_secret,
                access_token, access_token_secret)
    print("Keys saved to", filename)

def rate_limit(f):
    """
    A decorator to handle rate limiting from the Twitter API. If
    a rate limit error is encountered we will sleep until we can
    issue the API call again.
    """
    def new_f(*args, **kwargs):
        while True:
            resp = f(*args, **kwargs)
            if resp.status_code == 200:
                return resp
            elif resp.status_code == 429:
                reset = int(resp.headers['x-rate-limit-reset'])
                now = time.time()
                seconds = reset - now + 10
                if seconds < 1:
                    seconds = 10
                logging.warn("rate limit exceeded: sleeping %s secs", seconds)
                time.sleep(seconds)
            elif resp.status_code == 503:
                seconds = 60
                logging.warn("503 from Twitter API, sleeping %s", seconds)
                time.sleep(seconds)
            else:
                resp.raise_for_status()
    return new_f

class Twarc(object):
    """
    Your friendly neighborhood Twitter archiving class. Twarc allows
    you to search for existing tweets, stream live tweets that match
    a filter query and lookup (hdyrate) a list of tweet ids.

    Each method search, stream and hydrate returns a tweet iterator which
    allows you to do what you want with the data. Twarc handles rate limiting
    in the API, so it will go to sleep when Twitter tells it to, and wake back
    up when it is able to get more data from the API.
    """

    def __init__(self, consumer_key, consumer_secret, access_token,
                 access_token_secret):
        """
        Instantiate a Twarc instance. Make sure your environment variables
        are set.
        """

        self.consumer_key = consumer_key
        self.consumer_secret = consumer_secret
        self.access_token = access_token
        self.access_token_secret = access_token_secret
        self._connect()

    def search(self, q, max_id=None, since_id=None):
        """
        Pass in a query with optional max_id and min_id and get back
        an iterator for decoded tweets.
        """
        logging.info("starting search for %s", q)
        url = "https://api.twitter.com/1.1/search/tweets.json"
        params = {
            "count": 100,
            "q": q
        }

        while True:
            if since_id:
                params['since_id'] = since_id
            if max_id:
                params['max_id'] = max_id

            resp = self.get(url, params=params)
            statuses = resp.json()["statuses"]

            if len(statuses) == 0:
                logging.info("no new tweets matching %s", params)
                break

            for status in statuses:
                yield status

            max_id = str(int(status["id_str"]) - 1)

    def stream(self, query):
        """
        Returns an iterator for tweets that match a given filter query from
        the livestream of tweets happening right now.
        """
        url = 'https://stream.twitter.com/1.1/statuses/filter.json'
        params = {"track": query}
        headers = {'accept-encoding': 'deflate, gzip'}
        errors = 0
        while True:
            try:
                logging.info("connecting to filter stream for %s", query)
                resp = self.post(url, params, headers=headers, stream=True)
                errors = 0
                for line in resp.iter_lines(chunk_size=512):
                    try:
                        yield json.loads(line.decode())
                    except Exception as e:
                        logging.error("json parse error: %s - %s", e, line)
            except requests.exceptions.HTTPError as e:
                errors += 1
                logging.error(e)
                if e.response.status_code == 420:
                    t = errors * 60
                    logging.info("sleeping %s", t)
                    time.sleep(t)
                else:
                    t = errors * 5
                    logging.info("sleeping %s", t)
                    time.sleep(t)
            except Exception as e:
                errors += 1
                t = errors * 1
                logging.error(e)
                logging.info("sleeping %s", t)
                time.sleep(t)

    def hydrate(self, iterator):
        """
        Pass in an iterator of tweet ids and get back an iterator for the
        decoded JSON for each corresponding tweet.
        """
        ids = []
        url = "https://api.twitter.com/1.1/statuses/lookup.json"

        # lookup 100 tweets at a time
        for tweet_id in iterator:
            tweet_id = tweet_id.strip()  # remove new line if present
            ids.append(tweet_id)
            if len(ids) == 100:
                logging.info("hydrating %s ids", len(ids))
                resp = self.post(url, data={"id": ','.join(ids)})
                tweets = resp.json()
                tweets.sort(key=lambda t: t['id_str'])
                for tweet in tweets:
                    yield tweet
                ids = []

        # hydrate any remaining ones
        if len(ids) > 0:
            logging.info("hydrating %s", ids)
            resp = self.client.post(url, data={"id": ','.join(ids)})
            for tweet in resp.json():
                yield tweet

    @rate_limit
    def get(self, *args, **kwargs):
        try:
            return self.client.get(*args, **kwargs)
        except requests.exceptions.ConnectionError as e:
            logging.error("caught connection error %s", e)
            self._connect()
            return self.get(*args, **kwargs)

    @rate_limit
    def post(self, *args, **kwargs):
        try:
            return self.client.post(*args, **kwargs)
        except requests.exceptions.ConnectionError as e:
            logging.error("caught connection error %s", e)
            self._connect()
            return self.post(*args, **kwargs)

    def _connect(self):
        logging.info("creating http session")
        self.client = OAuth1Session(
            client_key=self.consumer_key,
            client_secret=self.consumer_secret,
            resource_owner_key=self.access_token,
            resource_owner_secret=self.access_token_secret
        )

if __name__ == "__main__":
    main()
<commit_msg>fixes #74<commit_after>#!/usr/bin/env python

from __future__ import print_function

import io
import os
import sys
import json
import time
import logging
import argparse
import requests

from requests_oauthlib import OAuth1Session

try:
    import configparser  # Python 3
except ImportError:
    import ConfigParser as configparser  # Python 2

if sys.version_info[:2] <= (2, 7):
    # Python 2
    get_input = raw_input
else:
    # Python 3
    get_input = input

def main():
    """
    The twarc command line.
    """
    parser = argparse.ArgumentParser("twarc")
    parser.add_argument("--search", dest="search",
                        help="search for tweets matching a query")
    parser.add_argument("--max_id", dest="max_id",
                        help="maximum tweet id to search for")
    parser.add_argument("--since_id", dest="since_id",
                        help="smallest id to search for")
    parser.add_argument("--stream", dest="stream",
                        help="stream tweets matching filter")
    parser.add_argument("--hydrate", dest="hydrate",
                        help="rehydrate tweets from a file of tweet ids")
    parser.add_argument("--log", dest="log",
                        default="twarc.log", help="log file")
    parser.add_argument("--consumer_key",
                        default=None, help="Twitter API consumer key")
    parser.add_argument("--consumer_secret",
                        default=None, help="Twitter API consumer secret")
    parser.add_argument("--access_token",
                        default=None, help="Twitter API access key")
    parser.add_argument("--access_token_secret",
                        default=None, help="Twitter API access token secret")
    parser.add_argument('-c', '--config',
                        default=default_config_filename(),
                        help="Config file containing Twitter keys and secrets")
    parser.add_argument('-p', '--profile', default='main',
                        help="Name of a profile in your configuration file")
    args = parser.parse_args()

    logging.basicConfig(
        filename=args.log,
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s"
    )

    consumer_key = args.consumer_key or os.environ.get('CONSUMER_KEY')
    consumer_secret = args.consumer_secret or os.environ.get('CONSUMER_SECRET')
    access_token = args.access_token or os.environ.get('ACCESS_TOKEN')
    access_token_secret = args.access_token_secret or os.environ.get("ACCESS_TOKEN_SECRET")

    if not (consumer_key and consumer_secret and
            access_token and access_token_secret):
        credentials = load_config(args.config, args.profile)
        if credentials:
            consumer_key = credentials['consumer_key']
            consumer_secret = credentials['consumer_secret']
            access_token = credentials['access_token']
            access_token_secret = credentials['access_token_secret']
        else:
            print("Please enter Twitter authentication credentials")
            consumer_key = get_input('consumer key: ')
            consumer_secret = get_input('consumer secret: ')
            access_token = get_input('access_token: ')
            access_token_secret = get_input('access token secret: ')
            save_keys(args.profile, consumer_key, consumer_secret,
                      access_token, access_token_secret)

    t = Twarc(consumer_key=consumer_key,
              consumer_secret=consumer_secret,
              access_token=access_token,
              access_token_secret=access_token_secret)

    if args.search:
        tweets = t.search(
            args.search,
            since_id=args.since_id,
            max_id=args.max_id
        )
    elif args.stream:
        tweets = t.stream(args.stream)
    elif args.hydrate:
        tweets = t.hydrate(open(args.hydrate, 'rU'))
    else:
        raise argparse.ArgumentTypeError(
            "must supply one of: --search --stream or --hydrate")

    # iterate through the tweets and write them to stdout
    for tweet in tweets:
        if "id_str" in tweet:
            logging.info("archived %s", tweet["id_str"])
            print(json.dumps(tweet))

def load_config(filename, profile):
    if not os.path.isfile(filename):
        return None
    config = configparser.ConfigParser()
    config.read(filename)
    data = {}
    for key in ['access_token', 'access_token_secret', 'consumer_key', 'consumer_secret']:
        try:
            data[key] = config.get(profile, key)
        except configparser.NoSectionError:
            sys.exit("no such profile %s in %s" % (profile, filename))
        except configparser.NoOptionError:
            sys.exit("missing %s from profile %s in %s" % (key, profile, filename))
    return data

def save_config(filename, profile,
                consumer_key, consumer_secret,
                access_token, access_token_secret):
    config = configparser.ConfigParser()
    config.add_section(profile)
    config.set(profile, 'consumer_key', consumer_key)
    config.set(profile, 'consumer_secret', consumer_secret)
    config.set(profile, 'access_token', access_token)
    config.set(profile, 'access_token_secret', access_token_secret)
    with open(filename, 'w') as config_file:
        config.write(config_file)

def default_config_filename():
    """
    Return the default filename for storing Twitter keys.
    """
    home = os.path.expanduser("~")
    return os.path.join(home, ".twarc")

def save_keys(profile, consumer_key, consumer_secret,
              access_token, access_token_secret):
    """
    Save keys to ~/.twarc
    """
    filename = default_config_filename()
    save_config(filename, profile,
                consumer_key, consumer_secret,
                access_token, access_token_secret)
    print("Keys saved to", filename)

def rate_limit(f):
    """
    A decorator to handle rate limiting from the Twitter API. If
    a rate limit error is encountered we will sleep until we can
    issue the API call again.
    """
    def new_f(*args, **kwargs):
        while True:
            resp = f(*args, **kwargs)
            if resp.status_code == 200:
                return resp
            elif resp.status_code == 429:
                reset = int(resp.headers['x-rate-limit-reset'])
                now = time.time()
                seconds = reset - now + 10
                if seconds < 1:
                    seconds = 10
                logging.warn("rate limit exceeded: sleeping %s secs", seconds)
                time.sleep(seconds)
            elif resp.status_code == 503:
                seconds = 60
                logging.warn("503 from Twitter API, sleeping %s", seconds)
                time.sleep(seconds)
            else:
                resp.raise_for_status()
    return new_f

class Twarc(object):
    """
    Your friendly neighborhood Twitter archiving class. Twarc allows
    you to search for existing tweets, stream live tweets that match
    a filter query and lookup (hdyrate) a list of tweet ids.

    Each method search, stream and hydrate returns a tweet iterator which
    allows you to do what you want with the data. Twarc handles rate limiting
    in the API, so it will go to sleep when Twitter tells it to, and wake back
    up when it is able to get more data from the API.
    """

    def __init__(self, consumer_key, consumer_secret, access_token,
                 access_token_secret):
        """
        Instantiate a Twarc instance. Make sure your environment variables
        are set.
        """

        self.consumer_key = consumer_key
        self.consumer_secret = consumer_secret
        self.access_token = access_token
        self.access_token_secret = access_token_secret
        self._connect()

    def search(self, q, max_id=None, since_id=None):
        """
        Pass in a query with optional max_id and min_id and get back
        an iterator for decoded tweets.
        """
        logging.info("starting search for %s", q)
        url = "https://api.twitter.com/1.1/search/tweets.json"
        params = {
            "count": 100,
            "q": q
        }

        while True:
            if since_id:
                params['since_id'] = since_id
            if max_id:
                params['max_id'] = max_id

            resp = self.get(url, params=params)
            statuses = resp.json()["statuses"]

            if len(statuses) == 0:
                logging.info("no new tweets matching %s", params)
                break

            for status in statuses:
                yield status

            max_id = str(int(status["id_str"]) - 1)

    def stream(self, query):
        """
        Returns an iterator for tweets that match a given filter query from
        the livestream of tweets happening right now.
        """
        url = 'https://stream.twitter.com/1.1/statuses/filter.json'
        params = {"track": query}
        headers = {'accept-encoding': 'deflate, gzip'}
        errors = 0
        while True:
            try:
                logging.info("connecting to filter stream for %s", query)
                resp = self.post(url, params, headers=headers, stream=True)
                errors = 0
                for line in resp.iter_lines(chunk_size=512):
                    # ignore keep alives
                    if line == "":
                        continue
                    try:
                        yield json.loads(line.decode())
                    except Exception as e:
                        logging.error("json parse error: %s - %s", e, line)
            except requests.exceptions.HTTPError as e:
                errors += 1
                logging.error(e)
                if e.response.status_code == 420:
                    t = errors * 60
                    logging.info("sleeping %s", t)
                    time.sleep(t)
                else:
                    t = errors * 5
                    logging.info("sleeping %s", t)
                    time.sleep(t)
            except Exception as e:
                errors += 1
                t = errors * 1
                logging.error(e)
                logging.info("sleeping %s", t)
                time.sleep(t)

    def hydrate(self, iterator):
        """
        Pass in an iterator of tweet ids and get back an iterator for the
        decoded JSON for each corresponding tweet.
        """
        ids = []
        url = "https://api.twitter.com/1.1/statuses/lookup.json"

        # lookup 100 tweets at a time
        for tweet_id in iterator:
            tweet_id = tweet_id.strip()  # remove new line if present
            ids.append(tweet_id)
            if len(ids) == 100:
                logging.info("hydrating %s ids", len(ids))
                resp = self.post(url, data={"id": ','.join(ids)})
                tweets = resp.json()
                tweets.sort(key=lambda t: t['id_str'])
                for tweet in tweets:
                    yield tweet
                ids = []

        # hydrate any remaining ones
        if len(ids) > 0:
            logging.info("hydrating %s", ids)
            resp = self.client.post(url, data={"id": ','.join(ids)})
            for tweet in resp.json():
                yield tweet

    @rate_limit
    def get(self, *args, **kwargs):
        try:
            return self.client.get(*args, **kwargs)
        except requests.exceptions.ConnectionError as e:
            logging.error("caught connection error %s", e)
            self._connect()
            return self.get(*args, **kwargs)

    @rate_limit
    def post(self, *args, **kwargs):
        try:
            return self.client.post(*args, **kwargs)
        except requests.exceptions.ConnectionError as e:
            logging.error("caught connection error %s", e)
            self._connect()
            return self.post(*args, **kwargs)

    def _connect(self):
        logging.info("creating http session")
        self.client = OAuth1Session(
            client_key=self.consumer_key,
            client_secret=self.consumer_secret,
            resource_owner_key=self.access_token,
            resource_owner_secret=self.access_token_secret
        )

if __name__ == "__main__":
    main()
<|endoftext|>