def apply_padders(cls):
    """ Decorator for generating useful padding methods
    Loops through and generates methods like:
      Padding.push_1(Widget)
      Sets the left padding attribute by 1
      Padding.pull_24(Widget)
      Sets right padding attribute by 24.
      Padding.center_50(Widget)
      Provides center padding with a relative width of 50
    """
    padding_count = 100

    for i in range(1, padding_count):
        setattr(cls, 'push_{}'.format(i), partialmethod(_Padding, left=i))
        setattr(cls, 'pull_{}'.format(i), partialmethod(_Padding, right=i))
        setattr(cls, 'center_{}'.format(i),
                partialmethod(_Padding, align='center',
                              width=('relative', i)))
        setattr(cls, 'left_{}'.format(i),
                partialmethod(_Padding, align='left',
                              width=('relative', i)))
        setattr(cls, 'right_{}'.format(i),
                partialmethod(_Padding, align='right',
                              width=('relative', i)))
    return cls
Пример #2
0
    def _op(argname, arg=None, *, _pseudop=pseudop):
        # this is a temportary helper to shrink the declarations of
        # the simpler operators. These operators all either take no
        # arguments, or take a single constant argument. this helper
        # is deleted after the operators are all defined.

        if arg is None:
            return partialmethod(_pseudop, Pseudop[argname])
        else:
            return partialmethod(_pseudop, Pseudop[argname], arg)
Пример #3
0
 def test_serialize_functools_partialmethod(self):
     value = functools.partialmethod(datetime.timedelta, 1, seconds=2)
     result = self.serialize_round_trip(value)
     self.assertIsInstance(result, functools.partialmethod)
     self.assertEqual(result.func, value.func)
     self.assertEqual(result.args, value.args)
     self.assertEqual(result.keywords, value.keywords)
Пример #4
0
 def on(self, command, func=None, **kwargs):
     
     if func is None:
         return functools.partialmethod(command, **kwargs)
     
     plugin = structures.Plugin(func, command, **kwargs)
     self.register(plugin)
     
     return plugin
def apply_style_map(cls):
    """ Applies AttrMap attributes to Color class
    Eg:
      Color.frame_header(Text("I'm text in the Orange frame header"))
      Color.body(Text("Im text in wrapped with the body color"))
    """
    for k in STYLES:
        setattr(cls, k[0], partialmethod(AttrMap, attr_map=k[0]))
    return cls
Пример #6
0
def override(cls, name, method):
    """Injects a method into a class.

    This will inject the given method into the given class, overriding a member
    of the given name. The new method will receive the same arguments as the old
    variant, with the original member passed in as second argument right after self.
    If the class previously had no member of this name, this argument will be None.
    """
    setattr(cls, name, partialmethod(method, getattr(cls, name, None)))
    return method
Пример #7
0
    def __new__(cls, name, bases, attrs):
        newClass = super().__new__(cls, name, bases, attrs)

        if hasattr(newClass, 'add'):
            add = getattr(newClass, 'add')
            for m in HTTP_METHODS:
                method = partialmethod(add, methods=m)
                setattr(newClass, m, method)
                setattr(newClass, m.lower(), method)

        return newClass
Пример #8
0
def admininline_factory(Model, Inline):
    classname = "{}{}".format(Model.__name__, "AdminInline")
    class_attributes = {"model": Model, "formset": OneFormFormSet, "min_num": 1, "max_num": 1, "can_delete": False}
    #{"model": Model, "extra": 2, "max_num": 1, "formset": inlineformset_factory(Release, Model,
    #  formset=DebugBaseInlineFormSet, fields="__all__", max_num=1, validate_max=True)}

    if hasattr(Model, "form_callback"):
        SpecificFormSet = type("{}SpecificFormSet".format(Model.__name__), (OneFormFormSet,),
                               {"collecster_instance_callback": functools.partialmethod(specific_instance_callback, Model=Model)})
        class_attributes["formset"] = SpecificFormSet

    return type(classname, (Inline,), class_attributes)
Пример #9
0
def register_tests(test_class, method_name, test_func, exclude=None):
    """
    Dynamically create serializer tests to ensure that all registered
    serializers are automatically tested.
    """
    formats = [
        f for f in serializers.get_serializer_formats()
        if (not isinstance(serializers.get_serializer(f), serializers.BadSerializer) and
            f != 'geojson' and
            (exclude is None or f not in exclude))
    ]
    for format_ in formats:
        setattr(test_class, method_name % format_, partialmethod(test_func, format_))
Пример #10
0
    def _create_go_to_method(cls, path, class_name=None):
        go_to_method = Navigation.GoToMethodFactory(path, class_name)
        inst_method = six.create_unbound_method(go_to_method, Navigation)

        # TODO(e0ne): remove python2 support once all integration jobs
        # will be switched to python3.
        if six.PY3:
            def _go_to_page(self, path):
                return Navigation._go_to_page(self, path)

            wrapped_go_to = functools.partialmethod(_go_to_page, path)
            setattr(Navigation, inst_method.name, wrapped_go_to)
        else:
            setattr(Navigation, inst_method.name, inst_method)
Пример #11
0
 def __init__(cls, name, bases, dct):
     func_names_to_proxy = dct.get('_proxyfunc_func_set_') or set(['__setitem__', '__getitem__', '__delitem__',
                                                                   '__contains__', '__iter__', '__len__'])
     for attr in func_names_to_proxy:
         if not hasattr(cls, attr):
             if sys.version_info[0] > 2:
                 # Python 3+
                 func = functools.partialmethod(dct['_proxyfunc_'], attr)
                 setattr(cls, attr, func)
             else:
                 # Python 2
                 func = (lambda attr:
                         lambda self, *args, **kwargs: dct['_proxyfunc_'](self, attr, *args, **kwargs))(attr)
                 setattr(cls, attr, new.instancemethod(func, None, cls))
     return super(ProxyContainerMethodsMetaClass, cls).__init__(name, bases, dct)
Пример #12
0
    def calculate(self, indices: list, n=None, parallel=True, cores=0) -> MutExResult:
        """
        :param indices: A list of indices for which to test the MutEx. The indices refer the the background-data row-ids.
        :return: MutExResult
        """

        if not all([x in self.background.index for x in indices]):
            raise Exception("Not all indices found in background")

        target = self.background.loc[indices]

        coverage = target.apply(max).pipe(sum)
        observation_signal = target.apply(sum, axis=1)
        signal = sum(observation_signal)

        if n == None:
            n = self.permutations

        logging.info("running {} permutations".format(n))

        if not parallel:
            cores = 1
        pool = mp.Pool(processes=mp.cpu_count() if cores < 1 else cores)
        logging.info('permutation with {} cores'.format(pool._processes))
        partial_simul = partialmethod(self._one_permutation)
        simulated_results = pool.starmap(partial_simul.func,
                                         zip(itertools.repeat(coverage, n), itertools.repeat(observation_signal, n)))
        pool.close()  # we are not adding any more processes
        pool.join()  # tell it to wait until all threads are done before going on

        logging.info('calculate result')
        sim_coverages = [x[0] for x in simulated_results]
        higher_coverage = [x[1] for x in simulated_results]
        lower_coverage = [x[2] for x in simulated_results]

        return MutExResult(coverage=coverage, signal=signal,
                           higher_coverage_count=np.sum(higher_coverage),
                           lower_coverage_count=np.sum(lower_coverage), permutations=n,
                           mean_sim_coverage=np.mean(sim_coverages),
                           stdev_sim_coverage=np.std(sim_coverages),
                           sample_size=len(self.sample_weights),
                           items=indices
                           )
Пример #13
0
	def __new__(mcs, name, bases, members):
		"""
		Defines a new TestCase class.

		This intercepts all tests that are parametrised and duplicates them for
		each of the sets of parameters they need to be called with.
		:param name: The name of the TestCase class.
		:param bases: The superclasses of the TestCase class.
		:param members: The members of the TestCase class, including functions.
		:return: A new TestCase class, modified for parametrised tests.
		"""
		members_copy = dict(members)
		for member in members_copy:
			if callable(members_copy[member]) and hasattr(members_copy[member], "parameters"): #It's a function that's marked with the @parametrise annotation.
				for test_name, parameters in members_copy[member].parameters.items(): #Copy the function for each set of parameters.
					new_function = functools.partialmethod(members_copy[member], **parameters) #Fill in only the parameters. The rest is filled in at calling (such as "self").
					members[member + "_" + test_name] = new_function #Store the filled-in function along with the test name to make it unique.
				del members[member] #Delete the original parametrised function.
		return type.__new__(mcs, name, bases, members)
Пример #14
0
                return curnode

def get_pairs(self, names):
        t = []
        for n in names:
                fine = n
                n = n.lower()
                n = n.replace('-', '')
                if n == 'description':
                        n = 'desc'
                value = getattr(self, n)
                if isinstance(value, list):
                        value = ", ".join(value)
                t.append((fine, value))
        return t


Person = collections.namedtuple('Person', ['name', 'desc', 'profile', 'joke'])
Person.get_pairs = partialmethod(get_pairs, ['Name', 'Description', 'Profile', 'Joke'])

Role = collections.namedtuple('Role', ['name', 'desc'])
Role.get_pairs = partialmethod(get_pairs, ['Name', 'Description'])

Circle = collections.namedtuple('Circle', \
                ['name', 'desc', 'domain', 'coremembers', 'tree'])
Circle.get_pairs = partialmethod(get_pairs, ['Name', 'Description', 'Domain', 'Core-members'])

Circle.__str__ = lambda self: str(self.name)
Project.get_pairs = get_pairs = partialmethod(get_pairs, ['Owner', 'Name', 'Description'])

Пример #15
0
    c_reboot              = partialmethod(_api_cmd, 'devmgr', 'restart', _req_params=['mac']) 
    c_force_provision     = partialmethod(_api_cmd, 'devmgr', 'force-provision', _req_params=['mac']) 
    c_poe_power_cycle     = partialmethod(_api_cmd, 'devmgr', 'power-cycle', _req_params=['mac', 'port_idx']) 
    c_adopt               = partialmethod(_api_cmd, 'devmgr', 'adopt', _req_params=['mac']) 
    c_speedtest           = partialmethod(_api_cmd, 'devmgr', 'speedtest') 
    c_speedtest_status    = partialmethod(_api_cmd, 'devmgr', 'speedtest-status') 
    c_set_locate          = partialmethod(_api_cmd, 'devmgr', 'set-locate', _req_params=['mac']) 
    c_unset_locate        = partialmethod(_api_cmd, 'devmgr', 'unset-locate', _req_params=['mac']) 
    c_upgrade             = partialmethod(_api_cmd, 'devmgr', 'upgrade', _req_params=['mac']) 
    c_upgrade_external    = partialmethod(_api_cmd, 'devmgr', 'upgrade-external', _req_params=['mac', 'url']) 
    c_spectrum_scan       = partialmethod(_api_cmd, 'devmgr', 'spectrum-scan', _req_params=['mac']) 

    c_backups             = partialmethod(_api_cmd, 'backup', 'list-backups')
    c_delete_backup       = partialmethod(_api_cmd, 'backup', 'delete-backup', _req_params=['filename'])

    c_make_backup         = partialmethod(_api_cmd, 'system', 'backup')
    c_check_firmware      = partialmethod(_api_cmd, 'system', 'check-firmware-update')

    c_clear_dpi           = partialmethod(_api_cmd, 'stat', 'reset-dpi')


''' A little python magic to automatically add device_macs and list_device for all known device 
types into the UnifiSite object '''
for dev_id in set(( x['type'] for x in DEVICES.values())):
    setattr(UnifiSite, "{}_macs".format(dev_id), partialmethod(UnifiSite.mac_by_type, dev_id) )
    setattr(UnifiSite, "list_{}".format(dev_id), partialmethod(UnifiSite.list_by_type, dev_id) )




Пример #16
0
    #print("METHOD", name)
    method = self._get_special(name)
    if method is NotImplemented:
        return NotImplemented
    return method()

unary_special_method_names = (
    "__len__",
    "__neg__", "__pos__", "__abs__",
    "__invert__", "__complex__",
    "__int__", "__float__",
    "__round__",
)

for name in unary_special_method_names:
    m = partialmethod(silk_unary_method, name=name)
    setattr(SilkBase, name, m)


def silk_unary_method_optional(self, name):
    #print("METHOD", name)
    try:
        method = self._get_special(name)
    except AttributeError:
        return NotImplementedError
    return method()

unary_special_method_names_optional = (
    "__length_hint__", "__index__",
    "__iter__", "__reversed__", "__bool__"
)
Пример #17
0
 def specialize(cls, model_name, **kwargs):
     assert set(kwargs) == set(cls._required_kwargs)
     specialized_class = type(
         "{}{}".format(cls.__name__, model_name), (cls, ),
         {'__init__': partialmethod(cls.__init__, **kwargs)})
     return specialized_class
Пример #18
0
class BasicCompletion(cmd2.Cmd):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def do_flag_based(self, statement: cmd2.Statement):
        """Tab completes arguments based on a preceding flag using flag_based_complete
        -f, --food [completes food items]
        -s, --sport [completes sports]
        -p, --path [completes local file system paths]
        """
        self.poutput("Args: {}".format(statement.args))

    def complete_flag_based(self, text, line, begidx, endidx) -> List[str]:
        """Completion function for do_flag_based"""
        flag_dict = {
            # Tab complete food items after -f and --food flags in command line
            '-f': food_item_strs,
            '--food': food_item_strs,
            # Tab complete sport items after -s and --sport flags in command line
            '-s': sport_item_strs,
            '--sport': sport_item_strs,
            # Tab complete using path_complete function after -p and --path flags in command line
            '-p': self.path_complete,
            '--path': self.path_complete,
        }

        return self.flag_based_complete(text,
                                        line,
                                        begidx,
                                        endidx,
                                        flag_dict=flag_dict)

    def do_index_based(self, statement: cmd2.Statement):
        """Tab completes first 3 arguments using index_based_complete"""
        self.poutput("Args: {}".format(statement.args))

    def complete_index_based(self, text, line, begidx, endidx) -> List[str]:
        """Completion function for do_index_based"""
        index_dict = {
            1:
            food_item_strs,  # Tab complete food items at index 1 in command line
            2:
            sport_item_strs,  # Tab complete sport items at index 2 in command line
            3: self.
            path_complete,  # Tab complete using path_complete function at index 3 in command line
        }

        return self.index_based_complete(text,
                                         line,
                                         begidx,
                                         endidx,
                                         index_dict=index_dict)

    def do_delimiter_complete(self, statement: cmd2.Statement):
        """Tab completes files from a list using delimiter_complete"""
        self.poutput("Args: {}".format(statement.args))

    # Use a partialmethod to set arguments to delimiter_complete
    complete_delimiter_complete = functools.partialmethod(
        cmd2.Cmd.delimiter_complete, match_against=file_strs, delimiter='/')

    def do_raise_error(self, statement: cmd2.Statement):
        """Demonstrates effect of raising CompletionError"""
        self.poutput("Args: {}".format(statement.args))

    def complete_raise_error(self, text, line, begidx, endidx) -> List[str]:
        """
        CompletionErrors can be raised if an error occurs while tab completing.

        Example use cases
            - Reading a database to retrieve a tab completion data set failed
            - A previous command line argument that determines the data set being completed is invalid
        """
        raise cmd2.CompletionError("This is how a CompletionError behaves")
Пример #19
0
        return f"{{:0{width}.{r_digits}f}}"

    def _formatter_function(self, selector: FormatSelector, value: int) -> int | float:
        """Return formatted value for display."""
        value = value * self.multiplier / self.divisor
        if self.unit_of_measurement == 0:
            # Zigbee spec power unit is kW, but we show the value in W
            value_watt = value * 1000
            if value_watt < 100:
                return round(value_watt, 1)
            return round(value_watt)
        if selector == self.FormatSelector.SUMMATION:
            return self._summa_format.format(value).lstrip()
        return self._format_spec.format(value).lstrip()

    demand_formatter = partialmethod(_formatter_function, FormatSelector.DEMAND)
    summa_formatter = partialmethod(_formatter_function, FormatSelector.SUMMATION)


@registries.ZIGBEE_CHANNEL_REGISTRY.register(smartenergy.Prepayment.cluster_id)
class Prepayment(ZigbeeChannel):
    """Prepayment channel."""


@registries.ZIGBEE_CHANNEL_REGISTRY.register(smartenergy.Price.cluster_id)
class Price(ZigbeeChannel):
    """Price channel."""


@registries.ZIGBEE_CHANNEL_REGISTRY.register(smartenergy.Tunneling.cluster_id)
class Tunneling(ZigbeeChannel):
Пример #20
0
class BaseValue(Resolveable):
    constant: Optional[
        bool] = None  # Whether the value depends on independent variables
    name: Optional[str] = None

    def _operation(
        self,
        function: Callable[..., Union[float, bool]],
        other: Union[float, Value, Empty] = _empty,
        reverse: bool = False,
    ):
        if not isinstance(other, Resolveable):
            other = Value(other)
        if reverse:
            assert other.value is not _empty
            return Value(Operation(function, other, self))
        return Value(Operation(function, self, other))

    __neg__ = partialmethod(_operation, operator.neg)
    __pos__ = partialmethod(_operation, operator.pos)
    __abs__ = partialmethod(_operation, operator.abs)

    __add__ = partialmethod(_operation, operator.add)
    __floordiv__ = partialmethod(_operation, operator.floordiv)
    __mod__ = partialmethod(_operation, operator.mod)
    __mul__ = partialmethod(_operation, operator.mul)
    __pow__ = partialmethod(_operation, operator.pow)
    __sub__ = partialmethod(_operation, operator.sub)
    __truediv__ = partialmethod(_operation, operator.truediv)

    __radd__ = partialmethod(_operation, operator.add, reverse=True)
    __rfloordiv__ = partialmethod(_operation, operator.floordiv, reverse=True)
    __rmod__ = partialmethod(_operation, operator.mod, reverse=True)
    __rmul__ = partialmethod(_operation, operator.mul, reverse=True)
    __rpow__ = partialmethod(_operation, operator.pow, reverse=True)
    __rsub__ = partialmethod(_operation, operator.sub, reverse=True)
    __rtruediv__ = partialmethod(_operation, operator.truediv, reverse=True)

    # TODO: Decide whether the implementation of the below six methods is really what we want
    __lt__ = partialmethod(_operation, operator.lt)
    __le__ = partialmethod(_operation, operator.le)
    __eq__ = partialmethod(_operation, operator.eq)  # type: ignore
    __ne__ = partialmethod(_operation, operator.ne)  # type: ignore
    __ge__ = partialmethod(_operation, operator.ge)
    __gt__ = partialmethod(_operation, operator.gt)
Пример #21
0
def setup_generic_relations(model_actor, relation_methods):
    """
    注册过的 actor_model_class 实例可以如此操作:
        actor.METHOD(target) # 建立 target 的特定 RELATION
        actor.dont_METHOD(target) # 或 no_METHOD 取消 target 的特定 RELATION
        #
        actor.actor_relations.filter(...)
        actor.METHOD_relations(...) # AS actor 的特定 RELATION filter
        actor.has_METHOD(target) # 与 target 是否存在特定 RELATION
    target_model_instance 可以有如下操作:
        target.target_relations.filter(...)
        target.METHOD_related(...) # AS target 的特定 RELATION filter
        target.was_METHOD(actor) # 与 actor 是否存在特定 RELATION
    owner_model_instance 可以有如下操作:
        owner.owner_relations.filter(...)
        owner.METHOD_related(...) # AS owner 的特定 RELATION filter
    """
    from qrelation import models
    #
    rel = GenericRelation(models.Relation,
        content_type_field='actor_type',
        object_id_field='actor_id',
        related_query_name='relations_with_%s_as_%s' % (label(model_actor), 'actor'))
    rel.contribute_to_class(model_actor, 'actor_relations')
    #
    actor_type = get_contenttype(model_actor)
    for method, kwargs in relation_methods.items():
        # 'follow': {
        #     'relation': models.REL_USER_FOLLOW,
        #     'target': 'quser.User',
        # }
        model_target = validate(kwargs.pop('target'))
        if not hasattr(model_target, 'target_relations'):
            rel = GenericRelation(models.Relation,
                content_type_field='target_type',
                object_id_field='target_id',
                related_query_name='relations_with_%s_as_%s' % (label(model_target), 'target'))
            rel.contribute_to_class(model_target, 'target_relations')
        model_owner = kwargs.pop('owner', None)
        if model_owner:
            model_owner = validate(model_owner)
            owner_type = get_contenttype(model_owner)
        if model_owner and not hasattr(model_owner, 'owner_relations'):
            rel = GenericRelation(models.Relation,
                content_type_field='owner_type',
                object_id_field='owner_id',
                related_query_name='relations_with_%s_as_%s' % (label(model_owner), 'owner'))
            rel.contribute_to_class(model_owner, 'owner_relations')
        #
        target_type = get_contenttype(model_target)
        relation, owner_field = kwargs.pop('relation'), kwargs.pop('owner_field', None)
        # 建立 relation
        setattr(model_actor, method, functools.partialmethod(
            relate, relation=relation, target_type=target_type, model_target=model_target, actor_type=actor_type, owner_field=owner_field))
        cancel_method = functools.partialmethod(
            relate, relation=relation, target_type=target_type, model_target=model_target, actor_type=actor_type, deleted=True)
        # 取消 relation
        setattr(model_actor, 'dont_%s' % method, cancel_method)
        setattr(model_actor, 'no_%s' % method, cancel_method)
        # 给 actor_model 增加类方法 METHOD_relations
        setattr(model_actor, '%s_relations' % method, functools.partialmethod(
            filter_relations, relation=relation))
        # 给 actor_model 增加类方法 has_METHOD
        setattr(model_actor, 'has_%s' % method, functools.partialmethod(
            check_relation, relation=relation, target_type=target_type, model_target=model_target))
        # 给 target_model 增加类方法 METHOD_related
        setattr(model_target, '%s_related' % method, functools.partialmethod(
            filter_related, relation=relation))
        # 给 target_model 增加类方法 was_METHOD
        setattr(model_target, 'was_%s' % method, functools.partialmethod(
            check_related, relation=relation, actor_type=actor_type, model_actor=model_actor))
        if not model_owner:
            continue
        # 给 owner_model 增加类方法 
        setattr(model_owner, '%s_related' % method, functools.partialmethod(
            filter_owner, relation=relation))
Пример #22
0
from functools import partialmethod
from requests import Session

class T:
    def test(self, msg=''):
        print(msg)

    test = partialmethod(test, msg='4444')

Session.get = partialmethod(Session.get, timeout=5)

T().test()
print(Session().get('http://www.baidu.com'))


from client import client
client._session.get('http://10.255.255.1')
 def _partialmethod(self):
     return partialmethod(self.func, *self.args, **self.keywords)
Пример #24
0
class APIClient(object):
    """
    An internal API client.

    >>> client = APIClient()
    >>> response = client.get('/projects/')
    >>> print response
    """
    def dispatch(
        self,
        path: str,
        method: str,
        data: dict = None,
        files: Mapping[str, BinaryIO] = None,
        json: dict = None,
        params: dict = None,
        request=None,
        tenant=True,
    ) -> Response:
        if request:
            assert not json
            assert not data
            assert not files
            data = request.data
            files = request.files
            json = None
            params = request.args

        if tenant is True:
            tenant = auth.get_current_tenant()

        if json:
            assert not data
            data = dumps(json, default=json_encoder)
            content_type = "application/json"
        elif files:
            if not data:
                data = {}
            for key, value in request.form.items():
                data[key] = value
            for key, value in files.items():
                data[key] = value
            content_type = "multipart/form-data"
        else:
            content_type = None

        with current_app.test_client() as client:
            response = client.open(
                path="/api/{}".format(path.lstrip("/")),
                query_string=params,
                method=method,
                content_type=content_type,
                data=data,
                environ_overrides={"zeus.tenant": tenant},
            )
        if not (200 <= response.status_code < 300):
            raise ApiError(text=response.get_data(as_text=True),
                           code=response.status_code)

        if response.headers["Content-Type"] != "application/json":
            raise ApiError(
                text="Request returned invalid content type: {}".format(
                    response.headers["Content-Type"]),
                code=response.status_code,
            )

        return response

    delete = partialmethod(dispatch, method="DELETE")
    get = partialmethod(dispatch, method="GET")
    head = partialmethod(dispatch, method="HEAD")
    options = partialmethod(dispatch, method="OPTIONS")
    patch = partialmethod(dispatch, method="PATCH")
    post = partialmethod(dispatch, method="POST")
    put = partialmethod(dispatch, method="PUT")
Пример #25
0
class Cluster(util.ListenableMixin, util.LocalLogMixin, metaclass=Registry):
    """A cluster on an endpoint"""
    _registry = {}
    _registry_range = {}
    _server_command_idx = {}
    _client_command_idx = {}

    def __init__(self, endpoint, is_server=True):
        self._endpoint = endpoint
        self._attr_cache = {}
        self._listeners = {}
        if is_server:
            self._type = ClusterType.Server
        else:
            self._type = ClusterType.Client

    @classmethod
    def from_id(cls, endpoint, cluster_id, is_server=True):
        if cluster_id in cls._registry:
            return cls._registry[cluster_id](endpoint, is_server)
        else:
            for cluster_id_range, cluster in cls._registry_range.items():
                if cluster_id_range[0] <= cluster_id <= cluster_id_range[1]:
                    c = cluster(endpoint, is_server)
                    c.cluster_id = cluster_id
                    return c

        LOGGER.warning("Unknown cluster %s", cluster_id)
        c = cls(endpoint, is_server)
        c.cluster_id = cluster_id
        return c

    def deserialize(self, tsn, frame_type, is_reply, command_id, data):
        if frame_type == 1:
            # Cluster command
            if is_reply:
                commands = self.client_commands
            else:
                commands = self.server_commands

            try:
                schema = commands[command_id][1]
                is_reply = commands[command_id][2]
            except KeyError:
                LOGGER.warning("Unknown cluster-specific command %s",
                               command_id)
                return tsn, command_id + 256, is_reply, data

            # Bad hack to differentiate foundation vs cluster
            command_id = command_id + 256
        else:
            # General command
            try:
                schema = foundation.COMMANDS[command_id][1]
                is_reply = foundation.COMMANDS[command_id][2]
            except KeyError:
                LOGGER.warning("Unknown foundation command %s", command_id)
                return tsn, command_id, is_reply, data

        value, data = t.deserialize(data, schema)
        if data != b'':
            LOGGER.warning("Data remains after deserializing ZCL frame")

        return tsn, command_id, is_reply, value

    @util.retryable_request
    def request(self,
                general,
                command_id,
                schema,
                *args,
                manufacturer=None,
                expect_reply=True):
        optional = len(
            [s for s in schema if hasattr(s, 'optional') and s.optional])
        if len(schema) < len(args) or len(args) < len(schema) - optional:
            self.error("Schema and args lengths do not match in request")
            error = asyncio.Future()
            error.set_exception(
                ValueError(
                    "Wrong number of parameters for request, expected %d argument(s)"
                    % len(schema)))
            return error

        sequence = self._endpoint._device.application.get_sequence()
        if general:
            frame_control = 0x00
        else:
            frame_control = 0x01
        if manufacturer is not None:
            frame_control |= 0b0100
            manufacturer = manufacturer.to_bytes(2, 'little')
        else:
            manufacturer = b''
        data = bytes([frame_control]) + manufacturer + bytes(
            [sequence, command_id])
        data += t.serialize(args, schema)

        return self._endpoint.request(self.cluster_id,
                                      sequence,
                                      data,
                                      expect_reply=expect_reply,
                                      command_id=command_id)

    def reply(self, general, command_id, schema, *args, manufacturer=None):
        if len(schema) != len(args) and foundation.Status not in schema:
            self.debug("Schema and args lengths do not match in reply")

        sequence = self._endpoint._device.application.get_sequence()
        frame_control = 0b1000  # Cluster reply command
        if not general:
            frame_control |= 0x01
        if manufacturer is not None:
            frame_control |= 0b0100
            manufacturer = manufacturer.to_bytes(2, 'little')
        else:
            manufacturer = b''
        data = bytes([frame_control]) + manufacturer + bytes(
            [sequence, command_id])
        data += t.serialize(args, schema)

        return self._endpoint.reply(self.cluster_id, sequence, data)

    def handle_message(self, is_reply, tsn, command_id, args):
        if is_reply:
            self.debug("Unexpected ZCL reply 0x%04x: %s", command_id, args)
            return

        self.debug("ZCL request 0x%04x: %s", command_id, args)
        if command_id <= 0xff:
            self.listener_event('zdo_command', tsn, command_id, args)
        else:
            # Unencapsulate bad hack
            command_id -= 256
            self.listener_event('cluster_command', tsn, command_id, args)
            self.handle_cluster_request(tsn, command_id, args)
            return

        if command_id == 0x0a:  # Report attributes
            valuestr = ", ".join([
                "%s=%s" %
                (self.attributes.get(a.attrid, [a.attrid])[0], a.value.value)
                for a in args[0]
            ])
            self.debug("Attribute report received: %s", valuestr)
            for attr in args[0]:
                self._update_attribute(attr.attrid, attr.value.value)
        else:
            self.handle_cluster_general_request(tsn, command_id, args)

    def handle_cluster_request(self, tsn, command_id, args):
        self.debug("No handler for cluster command %s", command_id)

    def handle_cluster_general_request(self, tsn, command_id, args):
        self.debug("No handler for general command %s", command_id)

    async def read_attributes_raw(self, attributes, manufacturer=None):
        schema = foundation.COMMANDS[0x00][1]
        attributes = [t.uint16_t(a) for a in attributes]
        v = await self.request(True,
                               0x00,
                               schema,
                               attributes,
                               manufacturer=manufacturer)
        return v

    async def read_attributes(self,
                              attributes,
                              allow_cache=False,
                              only_cache=False,
                              raw=False,
                              manufacturer=None):
        if raw:
            assert len(attributes) == 1
        success, failure = {}, {}
        attribute_ids = []
        orig_attributes = {}
        for attribute in attributes:
            if isinstance(attribute, str):
                attrid = self._attridx[attribute]
            else:
                attrid = attribute
            attribute_ids.append(attrid)
            orig_attributes[attrid] = attribute

        to_read = []
        if allow_cache or only_cache:
            for idx, attribute in enumerate(attribute_ids):
                if attribute in self._attr_cache:
                    success[attributes[idx]] = self._attr_cache[attribute]
                else:
                    to_read.append(attribute)
        else:
            to_read = attribute_ids

        if not to_read or only_cache:
            if raw:
                return success[attributes[0]]
            return success, failure

        result = await self.read_attributes_raw(to_read,
                                                manufacturer=manufacturer)
        if not isinstance(result[0], list):
            for attrid in to_read:
                orig_attribute = orig_attributes[attrid]
                failure[orig_attribute] = result[0]  # Assume default response
        else:
            for record in result[0]:
                orig_attribute = orig_attributes[record.attrid]
                if record.status == 0:
                    self._update_attribute(record.attrid, record.value.value)
                    success[orig_attribute] = record.value.value
                else:
                    failure[orig_attribute] = record.status

        if raw:
            # KeyError is an appropriate exception here, I think.
            return success[attributes[0]]
        return success, failure

    def write_attributes(self,
                         attributes,
                         is_report=False,
                         manufacturer=None,
                         unsupported_attrs=[]):
        args = []
        for attrid, value in attributes.items():
            if isinstance(attrid, str):
                attrid = self._attridx[attrid]
            if attrid not in self.attributes:
                self.error("%d is not a valid attribute id", attrid)
                continue

            if is_report:
                a = foundation.ReadAttributeRecord()
                a.status = 0
            else:
                a = foundation.Attribute()

            a.attrid = t.uint16_t(attrid)
            a.value = foundation.TypeValue()

            try:
                python_type = self.attributes[attrid][1]
                a.value.type = t.uint8_t(foundation.DATA_TYPE_IDX[python_type])
                a.value.value = python_type(value)
                args.append(a)
            except ValueError as e:
                self.error(str(e))

        if is_report and unsupported_attrs:
            for attrid in unsupported_attrs:
                a = foundation.ReadAttributeRecord()
                a.attrid = attrid
                a.status = foundation.Status.UNSUPPORTED_ATTRIBUTE
                args.append(a)

        if is_report:
            schema = foundation.COMMANDS[0x01][1]
            return self.reply(True,
                              0x01,
                              schema,
                              args,
                              manufacturer=manufacturer)
        else:
            schema = foundation.COMMANDS[0x02][1]
            return self.request(True,
                                0x02,
                                schema,
                                args,
                                manufacturer=manufacturer)

    def bind(self):
        return self._endpoint.device.zdo.bind(self._endpoint.endpoint_id,
                                              self.cluster_id)

    def unbind(self):
        return self._endpoint.device.zdo.unbind(self._endpoint.endpoint_id,
                                                self.cluster_id)

    def configure_reporting(self,
                            attribute,
                            min_interval,
                            max_interval,
                            reportable_change,
                            manufacturer=None):
        if isinstance(attribute, str):
            attrid = self._attridx.get(attribute, None)
        else:
            attrid = attribute
        if attrid not in self.attributes or attrid is None:
            self.error("{} is not a valid attribute id".format(attribute))
            return

        schema = foundation.COMMANDS[0x06][1]
        cfg = foundation.AttributeReportingConfig()
        cfg.direction = 0
        cfg.attrid = attrid
        cfg.datatype = foundation.DATA_TYPE_IDX.get(
            self.attributes.get(attrid, (None, None))[1], None)
        cfg.min_interval = min_interval
        cfg.max_interval = max_interval
        cfg.reportable_change = reportable_change
        return self.request(True,
                            0x06,
                            schema, [cfg],
                            manufacturer=manufacturer)

    def command(self, command, *args, manufacturer=None, expect_reply=True):
        schema = self.server_commands[command][1]
        return self.request(False,
                            command,
                            schema,
                            *args,
                            manufacturer=manufacturer,
                            expect_reply=expect_reply)

    def client_command(self, command, *args):
        schema = self.client_commands[command][1]
        return self.reply(False, command, schema, *args)

    @property
    def is_client(self) -> bool:
        """Return True if this is a client cluster."""
        return self._type == ClusterType.Client

    @property
    def is_server(self) -> bool:
        """Return True if this is a server cluster."""
        return self._type == ClusterType.Server

    @property
    def name(self):
        return self.__class__.__name__

    @property
    def endpoint(self):
        return self._endpoint

    @property
    def commands(self):
        return list(self._server_command_idx.keys())

    def _update_attribute(self, attrid, value):
        self._attr_cache[attrid] = value
        self.listener_event('attribute_updated', attrid, value)

    def log(self, lvl, msg, *args):
        msg = '[0x%04x:%s:0x%04x] ' + msg
        args = (
            self._endpoint.device.nwk,
            self._endpoint.endpoint_id,
            self.cluster_id,
        ) + args
        return LOGGER.log(lvl, msg, *args)

    def __getattr__(self, name):
        if name in self._client_command_idx:
            return functools.partial(
                self.client_command,
                self._client_command_idx[name],
            )
        elif name in self._server_command_idx:
            return functools.partial(
                self.command,
                self._server_command_idx[name],
            )
        else:
            raise AttributeError("No such command name: %s" % (name, ))

    def __getitem__(self, key):
        return self.read_attributes([key], allow_cache=True, raw=True)

    @util.retryable_request
    def _discover(self,
                  cmd_id,
                  start_item,
                  num_of_items,
                  manufacturer=None,
                  tries=3):
        schema = foundation.COMMANDS[cmd_id][1]
        return self.request(True,
                            cmd_id,
                            schema,
                            start_item,
                            num_of_items,
                            manufacturer=manufacturer)

    discover_attributes = functools.partialmethod(_discover, 0x0c)
    discover_attributes_extended = functools.partialmethod(_discover, 0x15)
    discover_commands_received = functools.partialmethod(_discover, 0x11)
    discover_commands_generated = functools.partialmethod(_discover, 0x13)
Пример #26
0
class Source:
    def __init__(self, entry=None, is_user=False):
        self.setters = dict()
        self.validate = list()
        self.is_user = is_user
        if entry:
            self.add_entry(entry, False)

    def add_setter(self, name, value, is_optional, conditions):
        self.setters.setdefault(name, list()).append(
            Setter(name, value, is_optional, conditions, self.is_user))

    def add_conditional_scope(self, scope, conditions):
        # TODO(cmaloney): 'defaults' are the same as 'can' and 'must' is identical to 'arguments' except
        # that one takes functions and one takes strings. Simplify to just 'can', 'must'.
        assert scope.keys() <= {'validate', 'default', 'must', 'conditional'}

        self.validate += scope.get('validate', list())

        for name, fn in scope.get('must', dict()).items():
            self.add_setter(name, fn, False, conditions)

        for name, fn in scope.get('default', dict()).items():
            self.add_setter(name, fn, True, conditions)

        for name, cond_options in scope.get('conditional', dict()).items():
            for value, sub_scope in cond_options.items():
                self.add_conditional_scope(sub_scope,
                                           conditions + [(name, value)])

    add_must = partialmethod(add_setter, is_optional=False, conditions=[])

    def add_value_dict(self, value_dict):
        for name, value in value_dict.items():
            self.add_must(name, value)

    def remove_setters(self, scope):
        def del_setter(name):
            if name in self.setters:
                del self.setters[name]

        for name in scope.get('must', dict()).keys():
            del_setter(name)

        for name in scope.get('default', dict()).keys():
            del_setter(name)

        for name, cond_options in scope.get('conditional', dict()).items():
            if name in self.setters:
                raise NotImplementedError(
                    "Should conditional setters overwrite all setters?")

    def add_entry(self, entry, replace_existing):
        if replace_existing:
            self.remove_setters(entry)

        self.add_conditional_scope(entry, [])

    def make_id(self):
        # {key: [hash_checkout(setter.make_id() for setter in setters)]
        #                 for key, setters in self.setters.items()}
        setter_ids = list()
        for setter_list in self.setters.values():
            for setter in setter_list:
                setter_ids.append(hash_checkout(setter.make_id()))
        return {
            'setters': setter_ids,
            'validate':
            [hash_checkout(function_id(fn)) for fn in self.validate],
            'is_user': self.is_user
        }
Пример #27
0
 class PartialClass(cls):
     __init__ = partialmethod(cls.__init__, *args, **kwargs)
Пример #28
0
 class PartialMeta(type):
     __call__ = functools.partialmethod(type(DictConverter).__call__,
                                        *expected_keys,
                                        delims=delims)
Пример #29
0
 class C(cl):
     __init__ = functools.partialmethod(cl.__init__, *args, **kwargs)
Пример #30
0
import argparse
import functools
import json
import logging
import platform
import re
import socket
import subprocess
import sys
import time
import urllib

logging.TRACE = 5
logging.addLevelName(logging.TRACE, 'TRACE')
logging.Logger.trace = functools.partialmethod(logging.Logger.log, logging.TRACE)
logging.trace = functools.partial(logging.log, logging.TRACE)
logging.basicConfig(format='%(message)s')
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

arch = platform.system()
# Windows does not support acsii colours
if not arch.startswith('Win'):
  WHITE = '\033[1m'
  RED = '\033[1;91m'
  GREEN = '\033[1;92m'
  YELLOW = '\033[1;93m'
  BLUE = '\033[1;94m'
  PURPLE = '\033[1;95m'
  NC = '\033[0m'
Пример #31
0
def mixed_scalar_binary_method_inplace(self, other, name2):
    result = getattr(self.value, name2)(other)
    if result is NotImplemented:
        return NotImplemented
    self._monitor.set_path(self._path, result)
    return result


class MixedScalar(MixedBase):
    def __getitem__(self, item):
        value = self.value
        if not isinstance(value, (str, np.ndarray)):
            raise TypeError(type(value))
        return value[item]


for name in binary_special_method_names:
    if name.startswith("__i"):
        name2 = "__" + name[3:]
        m = partialmethod(mixed_scalar_binary_method_inplace, name2=name2)
    else:
        m = partialmethod(mixed_scalar_binary_method, name=name)
    setattr(MixedScalar, name, m)

from .MixedDict import MixedDict, mixed_dict
from .MixedObject import MixedObject
from .Monitor import Monitor
from .MakeParentMonitor import MakeParentMonitor
from .OverlayMonitor import OverlayMonitor
from .get_form import is_contiguous, is_unsigned
            '\t {:<10} : {} ({})\n'.format('Pixel Type', self.pixeltype, self.dtype)+\
            '\t {:<10} : {}{}\n'.format('Components', self.components, ' (RGB)' if 'RGB' in self._libsuffix else '')+\
            '\t {:<10} : {}\n'.format('Dimensions', self.shape)+\
            '\t {:<10} : {}\n'.format('Spacing', tuple([round(s,4) for s in self.spacing]))+\
            '\t {:<10} : {}\n'.format('Origin', tuple([round(o,4) for o in self.origin]))+\
            '\t {:<10} : {}\n'.format('Direction', np.round(self.direction.flatten(),4))
        return s


if HAS_PY3:
    # Set partial class methods for any functions which take an ANTsImage as the first argument
    for k, v in utils.__dict__.items():
        if callable(v):
            args = inspect.getargspec(getattr(utils, k)).args
            if (len(args) > 0) and (args[0] in {'img', 'image'}):
                setattr(ANTsImage, k, partialmethod(v))

    for k, v in registration.__dict__.items():
        if callable(v):
            args = inspect.getargspec(getattr(registration, k)).args
            if (len(args) > 0) and (args[0] in {'img', 'image'}):
                setattr(ANTsImage, k, partialmethod(v))

    for k, v in segmentation.__dict__.items():
        if callable(v):
            args = inspect.getargspec(getattr(segmentation, k)).args
            if (len(args) > 0) and (args[0] in {'img', 'image'}):
                setattr(ANTsImage, k, partialmethod(v))

    for k, v in viz.__dict__.items():
        if callable(v):
Пример #33
0
        f.cl_ctx, f.cl_queue, f.ane, f.device = cl_ctx, cl_queue, ane, tt.device
        return f.apply(f, *x, **kwargs)

    setattr(Tensor, name, dispatch)
    if name in ['add', 'sub', 'mul', 'pow', 'matmul']:
        setattr(Tensor, f"__{name}__", dispatch)
        setattr(Tensor, f"__i{name}__",
                lambda self, x: self.assign(dispatch(self, x)))
        setattr(Tensor, f"__r{name}__", lambda self, x: dispatch(x, self))


for device in [
        device for device in Device.__dict__.keys() if device[0] != "_"
]:
    setattr(Tensor, f"{device.lower()}",
            functools.partialmethod(Tensor.to, Device.__dict__[device]))
    setattr(Tensor, f"{device.lower()}_",
            functools.partialmethod(Tensor.to_, Device.__dict__[device]))


# this registers all the operations
def _register_ops(namespace, device=Device.CPU):
    for name, cls in inspect.getmembers(namespace, inspect.isclass):
        if name[0] != "_": register(name.lower(), cls, device=device)


from tinygrad import ops_cpu
_register_ops(ops_cpu)
if os.getenv("CHERRY", None) is not None:
    from extra import ops_cherry
    _register_ops(ops_cherry)
Пример #34
0
 class NewCls(cls):
     __init__ = partialmethod(cls.__init__, *args, **kwargs)
Пример #35
0
class EZSP:

    COMMANDS = COMMANDS
    EZSP_VERSION = 4

    def __init__(self, device_config: Dict):
        self._awaiting = {}
        self._config = device_config
        self._callbacks = {}
        self._ezsp_event = asyncio.Event()
        self._seq = 0
        self._gw = None
        self._ezsp_version = self.EZSP_VERSION
        self._awaiting = {}
        self.COMMANDS_BY_ID = {}
        for name, details in self.COMMANDS.items():
            self.COMMANDS_BY_ID[details[0]] = (name, details[1], details[2])

    async def connect(self) -> None:
        assert self._gw is None
        self._gw = await uart.connect(self._config, self)

    @classmethod
    async def probe(cls, device_config: Dict) -> bool:
        """Probe port for the device presence."""
        ezsp = cls(SCHEMA_DEVICE(device_config))
        try:
            await asyncio.wait_for(ezsp._probe(), timeout=PROBE_TIMEOUT)
            return True
        except (asyncio.TimeoutError, serial.SerialException,
                APIException) as exc:
            LOGGER.debug(
                "Unsuccessful radio probe of '%s' port",
                device_config[CONF_DEVICE_PATH],
                exc_info=exc,
            )
        finally:
            ezsp.close()

        return False

    async def _probe(self) -> None:
        """Open port and try sending a command"""
        await self.connect()
        await self.reset()
        self.close()

    def reconnect(self):
        """Reconnect using saved parameters."""
        LOGGER.debug(
            "Reconnecting %s serial port on %s bauds",
            self._config[CONF_DEVICE_PATH],
            self._config[CONF_DEVICE_BAUDRATE],
        )
        return self.connect()

    async def reset(self):
        LOGGER.debug("Resetting EZSP")
        self.stop_ezsp()
        for seq in self._awaiting:
            future = self._awaiting[seq][2]
            if not future.done():
                future.cancel()
        self._awaiting = {}
        self._callbacks = {}
        self._seq = 0
        await self._gw.reset()
        self.start_ezsp()

    async def version(self):
        ver, stack_type, stack_version = await self._command(
            "version", self.ezsp_version)
        if ver != self.ezsp_version:
            self._ezsp_version = ver
            await self._command("version", ver)
            LOGGER.debug("Switched to EZSP protocol version %d",
                         self.ezsp_version)
        LOGGER.info("EZSP Stack Type: %s, Stack Version: %s", stack_type,
                    stack_version)

    def close(self):
        self.stop_ezsp()
        if self._gw:
            self._gw.close()
            self._gw = None

    def _ezsp_frame(self, name, *args):
        c = self.COMMANDS[name]
        data = t.serialize(args, c[1])
        frame = [self._seq & 0xFF, 0, c[0]]  # Frame control. TODO.  # Frame ID
        if self.ezsp_version >= 5:
            frame.insert(1, 0xFF)  # Legacy Frame ID
            frame.insert(1, 0x00)  # Ext frame control. TODO.

        return bytes(frame) + data

    def _command(self, name, *args):
        LOGGER.debug("Send command %s: %s", name, args)
        if not self.is_ezsp_running:
            raise EzspError("EZSP is not running")

        data = self._ezsp_frame(name, *args)
        self._gw.data(data)
        c = self.COMMANDS[name]
        future = asyncio.Future()
        self._awaiting[self._seq] = (c[0], c[2], future)
        self._seq = (self._seq + 1) % 256
        return asyncio.wait_for(future, timeout=EZSP_CMD_TIMEOUT)

    async def _list_command(self, name, item_frames, completion_frame, spos,
                            *args):
        """Run a command, returning result callbacks as a list"""
        fut = asyncio.Future()
        results = []

        def cb(frame_name, response):
            if frame_name in item_frames:
                results.append(response)
            elif frame_name == completion_frame:
                fut.set_result(response)

        cbid = self.add_callback(cb)
        try:
            v = await self._command(name, *args)
            if v[0] != t.EmberStatus.SUCCESS:
                raise Exception(v)
            v = await fut
            if v[spos] != t.EmberStatus.SUCCESS:
                raise Exception(v)
        finally:
            self.remove_callback(cbid)

        return results

    startScan = functools.partialmethod(
        _list_command,
        "startScan",
        ["energyScanResultHandler", "networkFoundHandler"],
        "scanCompleteHandler",
        1,
    )
    pollForData = functools.partialmethod(_list_command, "pollForData",
                                          ["pollHandler"],
                                          "pollCompleteHandler", 0)
    zllStartScan = functools.partialmethod(
        _list_command,
        "zllStartScan",
        ["zllNetworkFoundHandler"],
        "zllScanCompleteHandler",
        0,
    )
    rf4ceDiscovery = functools.partialmethod(
        _list_command,
        "rf4ceDiscovery",
        ["rf4ceDiscoveryResponseHandler"],
        "rf4ceDiscoveryCompleteHandler",
        0,
    )

    def connection_lost(self, exc):
        """Lost serial connection."""
        LOGGER.debug("%s connection lost unexpectedly: %s",
                     self._config[CONF_DEVICE_PATH], exc)
        self.enter_failed_state("Serial connection loss: {}".format(exc))

    def enter_failed_state(self, error):
        """UART received error frame."""
        LOGGER.error(
            "NCP entered failed state. Requesting APP controller restart")
        self.close()
        self.handle_callback("_reset_controller_application", (error, ))

    async def formNetwork(self, parameters):  # noqa: N802
        fut = asyncio.Future()

        def cb(frame_name, response):
            nonlocal fut
            if frame_name == "stackStatusHandler":
                fut.set_result(response)

        self.add_callback(cb)
        v = await self._command("formNetwork", parameters)
        if v[0] != t.EmberStatus.SUCCESS:
            raise Exception("Failure forming network: %s" % (v, ))

        v = await fut
        if v[0] != t.EmberStatus.NETWORK_UP:
            raise Exception("Failure forming network: %s" % (v, ))

        return v

    def __getattr__(self, name):
        if name not in self.COMMANDS:
            raise AttributeError("%s not found in COMMANDS" % (name, ))

        return functools.partial(self._command, name)

    def frame_received(self, data):
        """Handle a received EZSP frame

        The protocol has taken care of UART specific framing etc, so we should
        just have EZSP application stuff here, with all escaping/stuffing and
        data randomization removed.
        """
        sequence, frame_id, data = data[0], data[2], data[3:]
        if frame_id == 0xFF:
            frame_id = 0
            if len(data) > 1:
                frame_id = data[1]
                data = data[2:]

        frame_name = self.COMMANDS_BY_ID[frame_id][0]
        LOGGER.debug(
            "Application frame %s (%s) received: %s",
            frame_id,
            frame_name,
            binascii.hexlify(data),
        )

        if sequence in self._awaiting:
            expected_id, schema, future = self._awaiting.pop(sequence)
            assert expected_id == frame_id
            result, data = t.deserialize(data, schema)
            try:
                future.set_result(result)
            except asyncio.InvalidStateError:
                LOGGER.debug(
                    "Error processing %s response. %s command timed out?",
                    sequence,
                    self.COMMANDS_BY_ID.get(expected_id, [expected_id])[0],
                )
        else:
            schema = self.COMMANDS_BY_ID[frame_id][2]
            frame_name = self.COMMANDS_BY_ID[frame_id][0]
            result, data = t.deserialize(data, schema)
            self.handle_callback(frame_name, result)

    def add_callback(self, cb):
        id_ = hash(cb)
        while id_ in self._callbacks:
            id_ += 1
        self._callbacks[id_] = cb
        return id_

    def remove_callback(self, id_):
        return self._callbacks.pop(id_)

    def handle_callback(self, *args):
        for callback_id, handler in self._callbacks.items():
            try:
                handler(*args)
            except Exception as e:
                LOGGER.exception("Exception running handler", exc_info=e)

    def start_ezsp(self):
        """Mark EZSP as running."""
        self._ezsp_event.set()

    def stop_ezsp(self):
        """Mark EZSP stopped."""
        self._ezsp_event.clear()

    @property
    def is_ezsp_running(self):
        """Return True if EZSP is running."""
        return self._ezsp_event.is_set()

    @property
    def ezsp_version(self):
        """Return protocol version."""
        return self._ezsp_version
Пример #36
0
class ModernMarchModel(Model):
    def __init__(
        self,
        num_humans=50,
        belief_dims=30,
        p_1=0.1,
        p_2=0.9,
        p_hp=0.1,
        p_hm=0.1,
    ):
        # reset random seeds prior to each iteration
        np.random.seed()
        random.seed()
        # save configuration
        self.conf = {
            "num_humans": num_humans,
            "belief_dims": belief_dims,
            "p_1": p_1,
            "p_2": p_2,
            "p_hp": p_hp,
            "p_hm": p_hm,
        }
        self.running = True
        self.schedule = BaseScheduler(self)
        # init environment and data collector
        self.init_env()
        self.init_dc()

    def get_config(self, param, *args):
        return self.conf[param]

    # necessary in order to satisfy data collector interface
    get_belief_dims = partialmethod(get_config, "belief_dims")
    get_num_humans = partialmethod(get_config, "num_humans")
    get_p_1 = partialmethod(get_config, "p_1")
    get_p_2 = partialmethod(get_config, "p_2")
    get_p_hp = partialmethod(get_config, "p_hp")
    get_p_hm = partialmethod(get_config, "p_hm")

    def get_time(self, *args):
        return int(self.schedule.time)

    def init_env(self):
        # init reality
        r = Reality("R1", self)
        self.schedule.add(r)
        # init organization
        o = OrganizationalCode("O1", self)
        self.schedule.add(o)
        # init humans
        for i in range(self.conf["num_humans"]):
            h = Human(f"H{i+1}", self)
            self.schedule.add(h)
        return

    def init_dc(self):
        # data collector enables tracking of metric at each time step
        self.datacollector = DataCollector(
            model_reporters={
                "time": self.get_time,
                "belief_dims": self.get_belief_dims,
                "num_humans": self.get_num_humans,
                "p_1": self.get_p_1,
                "p_2": self.get_p_2,
                "p_hp": self.get_p_hp,
                "p_hm": self.get_p_hm,
                "code_kl": calc_code_kl,
                "human_kl": calc_human_kl,
                "human_kl_var": calc_kl_var,
                "human_kl_dissim": calc_dissim,
            })
        # collect metrics for time step 0
        self.datacollector.collect(self)
        return

    def get_exp_grp(self):
        # get list of humans with higher KL than code
        humans = self.schedule.agents[2:(2 + self.conf["num_humans"])]
        code = self.schedule.agents[1]
        return list(filter(lambda h: (h.kl > code.kl), humans))

    def step(self):
        try:
            # determine expert group for this time step
            self.exp_grp = self.get_exp_grp()
            # update all agents
            self.schedule.step()
            # collect metrics for this time step
            self.datacollector.collect(self)
        except Exception as e:
            # log potential erros, but continue with next iteration
            print("The following error occurred:")
            print(e)
            print("Model configuration:")
            print(self.conf)
Пример #37
0
class Git(Base):
    """Class for managing Git."""

    GITIGNORE = ".gitignore"
    GIT_DIR = ".git"
    LOCAL_BRANCH_PREFIX = "refs/heads/"
    RE_HEXSHA = re.compile(r"^[0-9A-Fa-f]{4,40}$")

    def __init__(self,
                 *args,
                 backends: Optional[Iterable[str]] = None,
                 **kwargs):
        self.ignored_paths: List[str] = []
        self.files_to_track: Set[str] = set()
        self.quiet: bool = False

        self.backends = GitBackends(backends, *args, **kwargs)
        first_ = first(self.backends.values())
        super().__init__(first_.root_dir)

    @property
    def dir(self):
        return first(self.backends.values()).dir

    @property
    def gitpython(self):
        return self.backends["gitpython"]

    @property
    def dulwich(self):
        return self.backends["dulwich"]

    @property
    def pygit2(self):
        return self.backends["pygit2"]

    @cached_property
    def stash(self):
        return Stash(self)

    @classmethod
    def clone(cls, url, to_path, **kwargs):
        for _, backend in GitBackends.DEFAULT.items():
            try:
                backend.clone(url, to_path, **kwargs)
                return Git(to_path)
            except NotImplementedError:
                pass
        raise NoGitBackendError("clone")

    @classmethod
    def is_sha(cls, rev):
        return rev and cls.RE_HEXSHA.search(rev)

    @staticmethod
    def _get_git_dir(root_dir):
        return os.path.join(root_dir, Git.GIT_DIR)

    @property
    def ignore_file(self):
        return self.GITIGNORE

    def _get_gitignore(self, path):
        ignore_file_dir = os.path.dirname(path)

        assert os.path.isabs(path)
        assert os.path.isabs(ignore_file_dir)

        entry = relpath(path, ignore_file_dir).replace(os.sep, "/")
        # NOTE: using '/' prefix to make path unambiguous
        if len(entry) > 0 and entry[0] != "/":
            entry = "/" + entry

        gitignore = os.path.join(ignore_file_dir, self.GITIGNORE)

        if not path_isin(os.path.realpath(gitignore), self.root_dir):
            raise FileNotInRepoError(path)

        return entry, gitignore

    def ignore(self, path):
        entry, gitignore = self._get_gitignore(path)

        if self.is_ignored(path):
            return

        msg = "Adding '{}' to '{}'.".format(relpath(path), relpath(gitignore))
        logger.debug(msg)

        self._add_entry_to_gitignore(entry, gitignore)

        self.track_file(relpath(gitignore))

        self.ignored_paths.append(path)

    def _add_entry_to_gitignore(self, entry, gitignore):
        entry = GitWildMatchPattern.escape(entry)

        with open(gitignore, "a+", encoding="utf-8") as fobj:
            unique_lines = set(fobj.readlines())
            fobj.seek(0, os.SEEK_END)
            if fobj.tell() == 0:
                # Empty file
                prefix = ""
            else:
                fobj.seek(fobj.tell() - 1, os.SEEK_SET)
                last = fobj.read(1)
                prefix = "" if last == "\n" else "\n"
            new_entry = f"{prefix}{entry}\n"
            if new_entry not in unique_lines:
                fobj.write(new_entry)

    def ignore_remove(self, path):
        entry, gitignore = self._get_gitignore(path)

        if not os.path.exists(gitignore):
            return

        with open(gitignore) as fobj:
            lines = fobj.readlines()

        filtered = list(filter(lambda x: x.strip() != entry.strip(), lines))

        if not filtered:
            os.unlink(gitignore)
            return

        with open(gitignore, "w") as fobj:
            fobj.writelines(filtered)

        self.track_file(relpath(gitignore))

    def _install_hook(self, name):
        hook = self._hook_path(name)
        with open(hook, "w+") as fobj:
            fobj.write(f"#!/bin/sh\nexec dvc git-hook {name} $@\n")

        os.chmod(hook, 0o777)

    def _install_merge_driver(self):
        self.gitpython.repo.git.config("merge.dvc.name", "DVC merge driver")
        self.gitpython.repo.git.config(
            "merge.dvc.driver",
            ("dvc git-hook merge-driver "
             "--ancestor %O "
             "--our %A "
             "--their %B "),
        )

    def install(self, use_pre_commit_tool=False):
        self._install_merge_driver()

        if not use_pre_commit_tool:
            self._verify_dvc_hooks()
            self._install_hook("post-checkout")
            self._install_hook("pre-commit")
            self._install_hook("pre-push")
            return

        from dvc.utils.serialize import modify_yaml

        config_path = os.path.join(self.root_dir, ".pre-commit-config.yaml")
        with modify_yaml(config_path) as config:
            entry = {
                "repo":
                "https://github.com/iterative/dvc",
                "rev":
                "master",
                "hooks": [
                    {
                        "id": "dvc-pre-commit",
                        "language_version": "python3",
                        "stages": ["commit"],
                    },
                    {
                        "id": "dvc-pre-push",
                        "language_version": "python3",
                        "stages": ["push"],
                    },
                    {
                        "id": "dvc-post-checkout",
                        "language_version": "python3",
                        "stages": ["post-checkout"],
                        "always_run": True,
                    },
                ],
            }

            if entry not in config["repos"]:
                config["repos"].append(entry)

    def cleanup_ignores(self):
        for path in self.ignored_paths:
            self.ignore_remove(path)
        self.reset_ignores()

    def reset_ignores(self):
        self.ignored_paths = []

    def reset_tracked_files(self):
        self.files_to_track = set()

    def remind_to_track(self):
        if self.quiet or not self.files_to_track:
            return

        files = " ".join(shlex.quote(path) for path in self.files_to_track)

        logger.info("\n"
                    "To track the changes with git, run:\n"
                    "\n"
                    "\tgit add {files}".format(files=files))

    def track_changed_files(self):
        if not self.files_to_track:
            return

        self.add(self.files_to_track)

    def track_file(self, path: str):
        self.files_to_track.add(path)

    def belongs_to_scm(self, path):
        basename = os.path.basename(path)
        path_parts = os.path.normpath(path).split(os.path.sep)
        return basename == self.ignore_file or Git.GIT_DIR in path_parts

    def has_rev(self, rev):
        try:
            self.resolve_rev(rev)
            return True
        except RevError:
            return False

    def close(self):
        self.backends.close_initialized()

    @cached_property
    def _hooks_home(self):
        return os.path.join(self.root_dir, self.GIT_DIR, "hooks")

    def _hook_path(self, name):
        return os.path.join(self._hooks_home, name)

    def _verify_hook(self, name):
        if os.path.exists(self._hook_path(name)):
            raise GitHookAlreadyExistsError(name)

    def _verify_dvc_hooks(self):
        self._verify_hook("post-checkout")
        self._verify_hook("pre-commit")
        self._verify_hook("pre-push")

    @property
    def no_commits(self):
        return not bool(self.get_ref("HEAD"))

    def _backend_func(self, name, *args, **kwargs):
        for backend in self.backends.values():
            try:
                func = getattr(backend, name)
                return func(*args, **kwargs)
            except NotImplementedError:
                pass
        raise NoGitBackendError(name)

    def get_fs(self, rev: str, **kwargs):
        from dvc.fs.git import GitFileSystem

        from .objects import GitTrie

        resolved = self.resolve_rev(rev)
        tree_obj = self.pygit2.get_tree_obj(rev=resolved)
        trie = GitTrie(tree_obj, resolved)
        return GitFileSystem(self.root_dir, trie, **kwargs)

    is_ignored = partialmethod(_backend_func, "is_ignored")
    add = partialmethod(_backend_func, "add")
    commit = partialmethod(_backend_func, "commit")
    checkout = partialmethod(_backend_func, "checkout")
    pull = partialmethod(_backend_func, "pull")
    push = partialmethod(_backend_func, "push")
    branch = partialmethod(_backend_func, "branch")
    tag = partialmethod(_backend_func, "tag")
    untracked_files = partialmethod(_backend_func, "untracked_files")
    is_tracked = partialmethod(_backend_func, "is_tracked")
    is_dirty = partialmethod(_backend_func, "is_dirty")
    active_branch = partialmethod(_backend_func, "active_branch")
    list_branches = partialmethod(_backend_func, "list_branches")
    list_tags = partialmethod(_backend_func, "list_tags")
    list_all_commits = partialmethod(_backend_func, "list_all_commits")
    get_rev = partialmethod(_backend_func, "get_rev")
    _resolve_rev = partialmethod(_backend_func, "resolve_rev")
    resolve_commit = partialmethod(_backend_func, "resolve_commit")

    set_ref = partialmethod(_backend_func, "set_ref")
    get_ref = partialmethod(_backend_func, "get_ref")
    remove_ref = partialmethod(_backend_func, "remove_ref")
    iter_refs = partialmethod(_backend_func, "iter_refs")
    iter_remote_refs = partialmethod(_backend_func, "iter_remote_refs")
    get_refs_containing = partialmethod(_backend_func, "get_refs_containing")
    push_refspec = partialmethod(_backend_func, "push_refspec")
    fetch_refspecs = partialmethod(_backend_func, "fetch_refspecs")
    _stash_iter = partialmethod(_backend_func, "_stash_iter")
    _stash_push = partialmethod(_backend_func, "_stash_push")
    _stash_apply = partialmethod(_backend_func, "_stash_apply")
    _stash_drop = partialmethod(_backend_func, "_stash_drop")
    describe = partialmethod(_backend_func, "describe")
    diff = partialmethod(_backend_func, "diff")
    reset = partialmethod(_backend_func, "reset")
    checkout_index = partialmethod(_backend_func, "checkout_index")
    status = partialmethod(_backend_func, "status")
    merge = partialmethod(_backend_func, "merge")

    def resolve_rev(self, rev: str) -> str:
        from dvc.repo.experiments.utils import exp_refs_by_name

        try:
            return self._resolve_rev(rev)
        except RevError:
            # backends will only resolve git branch and tag names,
            # if rev is not a sha it may be an abbreviated experiment name
            if not self.is_sha(rev) and not rev.startswith("refs/"):
                ref_infos = list(exp_refs_by_name(self, rev))
                if len(ref_infos) == 1:
                    return self.get_ref(str(ref_infos[0]))
                if len(ref_infos) > 1:
                    raise RevError(f"ambiguous Git revision '{rev}'")
            raise

    def branch_revs(self,
                    branch: str,
                    end_rev: Optional[str] = None) -> Iterable[str]:
        """Iterate over revisions in a given branch (from newest to oldest).

        If end_rev is set, iterator will stop when the specified revision is
        reached.
        """
        commit = self.resolve_commit(branch)
        while commit is not None:
            yield commit.hexsha
            parent = first(commit.parents)
            if parent is None or parent == end_rev:
                return
            commit = self.resolve_commit(parent)

    @contextmanager
    def detach_head(self, rev: Optional[str] = None):
        """Context manager for performing detached HEAD SCM operations.

        Detaches and restores HEAD similar to interactive git rebase.
        Restore is equivalent to 'reset --soft', meaning the caller is
        is responsible for preserving & restoring working tree state
        (i.e. via stash) when applicable.

        Yields revision of detached head.
        """
        if not rev:
            rev = "HEAD"
        orig_head = self.get_ref("HEAD", follow=False)
        logger.debug("Detaching HEAD at '%s'", rev)
        self.checkout(rev, detach=True)
        try:
            yield self.get_ref("HEAD")
        finally:
            prefix = self.LOCAL_BRANCH_PREFIX
            if orig_head.startswith(prefix):
                symbolic = True
                name = orig_head[len(prefix):]
            else:
                symbolic = False
                name = orig_head
            self.set_ref(
                "HEAD",
                orig_head,
                symbolic=symbolic,
                message=f"dvc: Restore HEAD to '{name}'",
            )
            logger.debug("Restore HEAD to '%s'", name)
            self.reset()

    @contextmanager
    def stash_workspace(self, **kwargs):
        """Stash and restore any workspace changes.

        Yields revision of the stash commit. Yields None if there were no
        changes to stash.
        """
        logger.debug("Stashing workspace")
        rev = self.stash.push(**kwargs)
        try:
            yield rev
        finally:
            if rev:
                logger.debug("Restoring stashed workspace")
                self.stash.pop()

    def _reset(self) -> None:
        self.backends.reset_all()
Пример #38
0
class DiscussionDetailHandler(base.OperationHandler):
    REPLIES_PER_PAGE = 50

    @base.require_perm(builtin.PERM_VIEW_DISCUSSION)
    @base.get_argument
    @base.route_argument
    @base.sanitize
    async def get(self, *, did: document.convert_doc_id, page: int = 1):
        ddoc = await discussion.inc_views(self.domain_id, did)
        if self.has_priv(builtin.PRIV_USER_PROFILE):
            dsdoc = await discussion.get_status(self.domain_id, ddoc['doc_id'],
                                                self.user['_id'])
        else:
            dsdoc = None
        vnode, (drdocs, pcount, drcount) = await asyncio.gather(
            discussion.get_vnode(self.domain_id, discussion.node_id(ddoc)),
            pagination.paginate(
                discussion.get_multi_reply(self.domain_id, ddoc['doc_id']),
                page, self.REPLIES_PER_PAGE))
        if not vnode:
            vnode = builtin.VNODE_MISSING
        elif vnode['doc_type'] == document.TYPE_PROBLEM and vnode.get(
                'hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        # TODO(twd2): do more visibility check eg. contest
        uids = {ddoc['owner_uid']}
        uids.update(drdoc['owner_uid'] for drdoc in drdocs)
        for drdoc in drdocs:
            if 'reply' in drdoc:
                uids.update(drrdoc['owner_uid'] for drrdoc in drdoc['reply'])
        if 'owner_uid' in vnode:
            uids.add(vnode['owner_uid'])
        udict, dudict = await asyncio.gather(
            user.get_dict(uids),
            domain.get_dict_user_by_uid(self.domain_id, uids))
        path_components = self.build_path(
            (self.translate('discussion_main'),
             self.reverse_url('discussion_main')),
            (vnode['title'],
             node_url(self, 'discussion_node', discussion.node_id(ddoc))),
            (ddoc['title'], None))
        self.render('discussion_detail.html',
                    page_title=ddoc['title'],
                    path_components=path_components,
                    ddoc=ddoc,
                    dsdoc=dsdoc,
                    drdocs=drdocs,
                    page=page,
                    pcount=pcount,
                    drcount=drcount,
                    vnode=vnode,
                    udict=udict,
                    dudict=dudict)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.require_perm(builtin.PERM_REPLY_DISCUSSION)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    @base.limit_rate('add_discussion', 3600, 30)
    async def post_reply(self, *, did: document.convert_doc_id, content: str):
        ddoc = await discussion.get(self.domain_id, did)
        await discussion.add_reply(self.domain_id, ddoc['doc_id'],
                                   self.user['_id'], content, self.remote_ip)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.require_perm(builtin.PERM_REPLY_DISCUSSION)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    @base.limit_rate('add_discussion', 3600, 30)
    async def post_tail_reply(self, *, did: document.convert_doc_id,
                              drid: document.convert_doc_id, content: str):
        ddoc = await discussion.get(self.domain_id, did)
        drdoc = await discussion.get_reply(self.domain_id, drid,
                                           ddoc['doc_id'])
        await discussion.add_tail_reply(self.domain_id, drdoc['doc_id'],
                                        self.user['_id'], content,
                                        self.remote_ip)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_edit_reply(self, *, did: document.convert_doc_id,
                              drid: document.convert_doc_id, content: str):
        ddoc = await discussion.get(self.domain_id, did)
        drdoc = await discussion.get_reply(self.domain_id, drid,
                                           ddoc['doc_id'])
        if (not self.own(
                ddoc, builtin.PERM_EDIT_DISCUSSION_REPLY_SELF_DISCUSSION) and
                not self.own(drdoc, builtin.PERM_EDIT_DISCUSSION_REPLY_SELF)):
            self.check_perm(builtin.PERM_EDIT_DISCUSSION_REPLY)
        drdoc = await discussion.edit_reply(self.domain_id,
                                            drdoc['doc_id'],
                                            content=content)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_delete_reply(self, *, did: document.convert_doc_id,
                                drid: document.convert_doc_id):
        ddoc = await discussion.get(self.domain_id, did)
        drdoc = await discussion.get_reply(self.domain_id, drid,
                                           ddoc['doc_id'])
        if (not self.own(ddoc,
                         builtin.PERM_DELETE_DISCUSSION_REPLY_SELF_DISCUSSION)
                and not self.own(drdoc,
                                 builtin.PERM_DELETE_DISCUSSION_REPLY_SELF)):
            self.check_perm(builtin.PERM_DELETE_DISCUSSION_REPLY)
        await oplog.add(self.user['_id'],
                        oplog.TYPE_DELETE_DOCUMENT,
                        doc=drdoc)
        drdoc = await discussion.delete_reply(self.domain_id, drdoc['doc_id'])
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_edit_tail_reply(self, *, did: document.convert_doc_id,
                                   drid: document.convert_doc_id,
                                   drrid: document.convert_doc_id,
                                   content: str):
        ddoc = await discussion.get(self.domain_id, did)
        drdoc, drrdoc = await discussion.get_tail_reply(
            self.domain_id, drid, drrid)
        if not drdoc or drdoc['parent_doc_id'] != ddoc['doc_id']:
            raise error.DocumentNotFoundError(domain_id,
                                              document.TYPE_DISCUSSION_REPLY,
                                              drid)
        if (not self.own(
                ddoc, builtin.PERM_EDIT_DISCUSSION_REPLY_SELF_DISCUSSION) and
                not self.own(drrdoc, builtin.PERM_EDIT_DISCUSSION_REPLY_SELF)):
            self.check_perm(builtin.PERM_EDIT_DISCUSSION_REPLY)
        await discussion.edit_tail_reply(self.domain_id, drid, drrid, content)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_delete_tail_reply(self, *, did: document.convert_doc_id,
                                     drid: document.convert_doc_id,
                                     drrid: objectid.ObjectId):
        ddoc = await discussion.get(self.domain_id, did)
        drdoc, drrdoc = await discussion.get_tail_reply(
            self.domain_id, drid, drrid)
        if not drdoc or drdoc['parent_doc_id'] != ddoc['doc_id']:
            raise error.DocumentNotFoundError(domain_id,
                                              document.TYPE_DISCUSSION_REPLY,
                                              drid)
        if (not self.own(ddoc,
                         builtin.PERM_DELETE_DISCUSSION_REPLY_SELF_DISCUSSION)
                and not self.own(drrdoc,
                                 builtin.PERM_DELETE_DISCUSSION_REPLY_SELF)):
            self.check_perm(builtin.PERM_DELETE_DISCUSSION_REPLY)
        await oplog.add(self.user['_id'],
                        oplog.TYPE_DELETE_SUB_DOCUMENT,
                        sub_doc=drrdoc,
                        doc_type=drdoc['doc_type'],
                        doc_id=drdoc['doc_id'])
        await discussion.delete_tail_reply(self.domain_id, drid, drrid)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.require_perm(builtin.PERM_VIEW_DISCUSSION)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def star_unstar(self, *, did: document.convert_doc_id, star: bool):
        ddoc = await discussion.get(self.domain_id, did)
        ddoc = await discussion.set_star(self.domain_id, ddoc['doc_id'],
                                         self.user['_id'], star)
        self.json_or_redirect(self.url, star=ddoc['star'])

    post_star = functools.partialmethod(star_unstar, star=True)
    post_unstar = functools.partialmethod(star_unstar, star=False)
Пример #39
0
class Moderation(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

        self.slowmodes = JSONFile('slowmodes.json')
        self.slowmode_bucket = {}

        if hasattr(self.bot, '__mod_mute_role_create_bucket__'):
            self._mute_role_create_cooldowns = self.bot.__mod_mute_role_create_bucket__
        else:
            self._mute_role_create_cooldowns = commands.CooldownMapping.from_cooldown(
                2, 600, commands.BucketType.guild)

    def cog_unload(self):
        self.bot.__mod_mute_role_create_bucket__ = self._mute_role_create_cooldowns

    async def call_mod_log_invoke(self, invoke, ctx):
        mod_log = ctx.bot.get_cog('ModLog')
        if mod_log:
            await getattr(mod_log, f'mod_{invoke}')(ctx)

    __before_invoke = functools.partialmethod(call_mod_log_invoke,
                                              'before_invoke')
    __after_invoke = functools.partialmethod(call_mod_log_invoke,
                                             'after_invoke')

    @staticmethod
    def _is_slowmode_immune(member):
        return member.guild_permissions.manage_guild

    async def check_slowmode(self, message):
        if not message.guild:
            return

        guild_id = message.guild.id
        if guild_id not in self.slowmodes:
            return

        slowmodes = self.slowmodes[guild_id]

        is_immune = self._is_slowmode_immune(message.author)

        for thing in (message.channel, message.author):
            key = str(thing.id)
            if key not in slowmodes:
                continue

            config = slowmodes[key]
            if not config['no_immune'] and is_immune:
                continue

            bucket = self.slowmode_bucket.setdefault(thing.id, {})
            time = bucket.get(message.author.id)
            if not time or (message.created_at -
                            time).total_seconds() >= config['duration']:
                bucket[message.author.id] = message.created_at
            else:
                await message.delete()
                break

    @commands.group(name='slowmode', invoke_without_command=True)
    @commands.has_permissions(manage_messages=True)
    @commands.bot_has_permissions(manage_messages=True)
    async def _slowmode(self,
                        ctx,
                        duration: time.Delta,
                        *,
                        member: discord.Member = None):
        """Activates the slowmode.

        If a member is given as an argument, it puts only the member in slowmode on the whole server,
        otherwise it puts the current channel in slowmode for all.

        Those with Manage Server permissions will not be affected. If you want to put them in slowmode too, use `{prefix}slowmode noimmune`.
        """

        pronoun = 'They'
        if not member:
            member = ctx.channel
            pronoun = 'Everyone'
        elif self._is_slowmode_immune(member):
            message = (
                f'{member} is immune from slowmode due to having the Manage Server permission. '
                f'Consider using `{ctx.prefix}slowmode noimmune`.')
            return await ctx.send(message)

        config = self.slowmodes.get(ctx.guild.id, {})
        slowmode = config.setdefault(str(member.id), {'no_immune': False})
        if slowmode['no_immune']:
            return await ctx.send(
                f'{member.mention} is already in **noimmune** slowmode. You need to turn it off first.'
            )

        slowmode['duration'] = duration.duration
        await self.slowmodes.put(ctx.guild.id, config)

        await ctx.send(
            f'{member.mention} is now in slowmode! {pronoun} must wait {duration} between each message they send.'
        )

    @_slowmode.command(name='noimmune')
    @commands.has_permissions(manage_messages=True)
    @commands.bot_has_permissions(manage_messages=True)
    async def _slowmode_no_immune(self,
                                  ctx,
                                  duration: time.Delta,
                                  *,
                                  member: discord.Member = None):
        """Puts the channel or member in "noimmune" slowmode.

        Unlike `{prefix}slowmode`, no one is immune to this slowmode, even those with Manage Server permissions.
        """

        if not member:
            member, pronoun = ctx.channel, 'They'
        else:
            pronoun = 'Everyone'

        config = self.slowmodes.get(ctx.guild.id, {})
        slowmode = config.setdefault(str(member.id), {'no_immune': True})
        slowmode['duration'] = duration.duration
        await self.slowmodes.put(ctx.guild, config)

        await ctx.send(
            f'{member.mention} is now in **noimmune** slowmode. {pronoun} must wait {duration} after each message they send.'
        )

    @_slowmode.command(name='off')
    async def _slowmode_off(self, ctx, *, member: discord.Member = None):
        """Turns off the slowmode for either a member or a channel."""

        member = member or ctx.channel
        config = self.slowmodes.get(ctx.guild.id, {})
        try:
            del config[str(member.id)]
        except ValueError:
            await ctx.send(f'{member.mention} was never in slowmode.')
        else:
            await self.slowmodes.put(ctx.guild.id, config)
            self.slowmode_bucket.pop(member.id, None)
            await ctx.send(f'{member.mention} is no longer in slowmode.')

    @commands.command(name='newusers', aliases=['newmembers', 'joined'])
    @commands.guild_only()
    async def _new_users(self, ctx, *, count: int = 5):
        """Tells you the recently joined members on this server.

        The minimum is 3 members. If no number is given, I will show the last 5 members that joined.
        """

        human_delta = time.human_timedelta
        count = max(count, 3)
        members = heapq.nlargest(count,
                                 ctx.guild.members,
                                 key=attrgetter('joined_at'))

        names = map(str, members)
        values = ((
            f'**Joined:** {human_delta(member.joined_at)}\n'
            f'**Created:** {human_delta(member.created_at)}\n{"-" * 40}')
                  for member in members)
        entries = zip(names, values)

        title = f'The {formats.pluralize(**{"newest members": len(members)})}'
        pages = FieldPaginator(ctx,
                               entries,
                               per_page=5,
                               title=title,
                               color=random_color())
        await pages.interact()

    @commands.command(name='clear')
    @commands.has_permissions(manage_messages=True)
    @commands.bot_has_permissions(manage_messages=True)
    async def _clear(self, ctx, arg: typing.Union[int, discord.Member]):
        """Clears some messages in a channel.

        The argument can either be a user or a number. If it's a number, it deletes *up to* that many messages.
        If it's a user, it deletes any message by that user up to the last 100 messages.
        If no argument was specified, it deletes my messages.
        """

        if isinstance(arg, int):
            if arg < 1:
                return await ctx.send(f'How can I delete {arg} messages?')

            deleted = await ctx.channel.purge(limit=min(arg, 1000) + 1)
        elif isinstance(arg, discord.Member):
            deleted = await ctx.channel.purge(
                check=lambda m: m.author.id == arg.id)

        messages = formats.pluralize(message=len(deleted) - 1)
        await ctx.send(f'Successfully deleted {messages}.', delete_after=2)

    @commands.command(name='clean')
    @commands.has_permissions(manage_messages=True)
    @commands.guild_only()
    async def _clean(self, ctx, limit=100):
        """Cleans up my messages from the channel.

        Give me Manage Messages and Read Message History permissions, and I'll also delete messages that invoked my commands.
        """

        prefixes = tuple(self.bot.get_guild_prefixes(ctx.guild))
        bot_id = self.bot.user.id

        bot_perms = ctx.channel.permissions_for(ctx.me)
        purge = functools.partial(ctx.channel.purge,
                                  limit=limit,
                                  before=ctx.message)
        can_bulk_delete = bot_perms.manage_messages and bot_perms.read_message_history

        if can_bulk_delete:

            def is_possible_command_invoke(m):
                if m.author.id == bot_id:
                    return True

                return m.content.startswith(
                    prefixes) and not m.content[1:2].isspace()

            deleted = await purge(check=is_possible_command_invoke)

        else:
            deleted = await purge(check=lambda m: m.author.id == bot_id,
                                  bulk=False)

        spammers = Counter(str(m.author) for m in deleted)
        total_deleted = sum(spammers.values())

        second_part = ' was' if total_deleted == 1 else 's were'
        title = f'{total_deleted} message{second_part} removed.'
        joined = '\n'.join(
            itertools.starmap('**{0}**: {1}'.format, spammers.most_common()))

        if ctx.bot_has_embed_links():
            spammer_stats = joined or discord.Embed.Empty

            embed = (discord.Embed(
                description=spammer_stats,
                color=random_color()).set_author(name=title))
            embed.timestamp = ctx.message.created_at

            await ctx.send(embed=embed, delete_after=10)

        else:
            message = f'{title}\n{joined}'
            await ctx.send(message, delete_after=10)

        await asyncio.sleep(20)
        with contextlib.suppress(discord.HTTPException):
            await ctx.message.delete()

    @_clear.error
    @_clean.error
    async def _clear_error(self, ctx, error):
        cause = error.__cause__
        if not isinstance(cause, discord.HTTPException):
            ctx.__bypass_local_error__ = True
            return

        await ctx.send(
            f'Couldn\'t delete the messages for some reason. Here\'s the error:\n```py\n{type(cause).__name__}: {cause}```'
        )

    @staticmethod
    async def _get_warn_timeout(connection, guild_id):
        query = 'SELECT timeout FROM warn_timeouts WHERE guild_id = $1;'
        row = await connection.fetchrow(query, guild_id)
        return row['timeout'] if row else datetime.timedelta(minutes=15)

    @commands.command(name='warn')
    @commands.has_permissions(manage_guild=True)
    async def _warn(self, ctx, member: discord.Member, *, reason):
        """Warns a user."""

        author, current_time, guild_id = ctx.author, ctx.message.created_at, ctx.guild.id
        timeout = await self._get_warn_timeout(ctx.db, guild_id)

        query = """
            SELECT   warned_at
            FROM     warn_entries
            WHERE    guild_id = $1 AND user_id = $2 AND warned_at + $3 > $4
            ORDER BY id;
        """
        records = await ctx.db.fetch(query, guild_id, member.id, timeout,
                                     current_time)
        warn_queue = [record[0] for record in records]

        try:
            last_warn = warn_queue[-1]
        except IndexError:
            pass
        else:
            retry_after = (current_time - last_warn).total_seconds()
            if retry_after <= 60:
                raise AlreadyWarned(
                    f'{member} has been warned already, try again in {60 - retry_after :.2f} seconds.'
                )

        query = """
            INSERT INTO warn_entries (guild_id, user_id, mod_id, reason, warned_at)
            VALUES      ($1, $2, $3, $4, $5);
        """
        await ctx.db.execute(query, guild_id, member.id, author.id, reason,
                             current_time)

        current_warn_number = len(warn_queue) + 1
        query = 'SELECT type, duration FROM warn_punishments WHERE guild_id = $1 AND warns = $2;'
        row = await ctx.db.fetchrow(query, guild_id, current_warn_number)

        if not row:
            if current_warn_number == 3:
                row = _default_punishment
            else:
                return await ctx.send(
                    f'\N{WARNING SIGN} Warned {member.mention} successfully!')

        # Auto-punish this f****t who dares to break the rules
        args = [member]
        duration = row['duration']
        if duration > 0:
            args.append(duration)
            punished_for = f' for {time.duration_units(duration)}'
        else:
            punished_for = ''

        punishment = row['type']
        punishment_command = getattr(self, punishment)
        punishment_reason = f'{reason}\n({ordinal(current_warn_number)} warning)'

        with temporary_attribute(ctx, 'send',
                                 lambda *a, **kw: asyncio.sleep(0)):
            await ctx.invoke(punishment_command,
                             *args,
                             reason=punishment_reason)

        message = (
            f'{member.mention} has {current_warn_number} warnings! Mate, you f****d up. Now take a {punishment}{punished_for}.'
        )
        await ctx.send(message)

        ctx.auto_punished = True
        ctx.command = punishment_command
        ctx.args[2:] = args
        ctx.kwargs['reason'] = punishment_reason

    @_warn.error
    async def _warn_error(self, ctx, error):
        if isinstance(error, AlreadyWarned):
            await ctx.send(error)
        else:
            ctx.__bypass_local_error__ = True

    @commands.command(name='warns')
    @commands.has_permissions(manage_guild=True)
    async def _warns(self, ctx, *, member: discord.Member = None):
        """Shows a given user's warns on this server.

        If no user is given, this command shows all warned users on this server.
        """

        if not member:
            query = 'SELECT user_id, reason FROM warn_entries WHERE guild_id = $1;'
            records = await ctx.db.fetch(query, ctx.guild.id)

            title = f'Warns in {ctx.guild}'
        else:
            query = 'SELECT user_id, reason FROM warn_entries WHERE guild_id = $1 AND user_id = $2;'
            records = await ctx.db.fetch(query, ctx.guild.id, member.id)

            title = f'Warns for {member}'

        entries = (itertools.starmap('`{0}.` <@{1[0]}> => **{1[1]}**'.format,
                                     enumerate(records, 1)) if records else
                   ('No warns found.', ))

        pages = Paginator(ctx, entries, per_page=5, title=title)
        await pages.interact()

    @commands.command(name='clearwarns', aliases=['resetwarns'])
    @commands.has_permissions(manage_guild=True)
    async def _clear_warns(self, ctx, member: discord.Member):
        """Clears a member's warns."""

        query = 'DELETE FROM warn_entries WHERE guild_id = $1 AND user_id = $2;'
        await ctx.db.execute(query, ctx.guild.id, member.id)

        await ctx.send(f'{member}\'s warns have been reset!')

    @commands.command(name='warnpunish')
    @commands.has_permissions(manage_guild=True)
    async def _warn_punish(self, ctx, num: num_warns, *,
                           punishment: warn_punishment):
        """Sets the punishment a user receives upon exceeding a given warn limit.

        Valid punishments:
            `mute` (Requires a given duration)
            `kick`
            `softban`
            `tempban` (Requires a given duration)
            `ban`
        """

        punishment, duration = punishment
        true_duration = None if not duration else duration.duration

        query = """
            INSERT INTO   warn_punishments (guild_id, warns, type, duration)
            VALUES        ($1, $2, $3, $4)
            ON CONFLICT   (guild_id, warns)
            DO UPDATE SET type = $3, duration = $4;
        """
        await ctx.db.execute(query, ctx.guild.id, num, punishment,
                             true_duration)

        extra = f' for {duration}' if duration else ''
        await ctx.send(
            f'\N{OK HAND SIGN} If a user has been warned {num} times, I will {punishment} them{extra}.'
        )

    @commands.command(name='warnpunishments')
    async def _warn_punishments(self, ctx):
        """Shows the list of warn punishments."""

        query = """
            SELECT   warns, initcap(type), duration
            FROM     warn_punishments
            WHERE    guild_id = $1
            ORDER BY warns;
        """
        punishments = await ctx.db.fetch(
            query, ctx.guild.id) or (_default_punishment, )

        entries = (
            f'{warns} strikes => **{type}** {(f"for " + time.duration_units(duration)) if duration else ""}'
            for warns, type, duration in punishments)

        pages = Paginator(ctx, entries, title=f'Punishments for {ctx.guild}')
        await pages.interact()

    @commands.command(name='warntimeout')
    @commands.has_permissions(manage_guild=True)
    async def _warn_timeout(self, ctx, duration: time.Delta):
        """Sets the maximum time between the oldest and the most recent warn.

        If a user hits a warn limit within this timeframe, they will be punished.
        """

        query = """
            INSERT INTO   warn_timeouts (guild_id, timeout)
            VALUES        ($1, $2)
            ON CONFLICT   (guild_id)
            DO UPDATE SET timeout = $2;
        """
        await ctx.db.execute(query, ctx.guild.id,
                             datetime.timedelta(seconds=duration.duration))

        await ctx.send(
            f'Aye, if a user was warned within **{duration}** after the oldest warn, bad things are going to happen.'
        )

    async def _get_muted_role_from_db(self, guild, *, connection=None):
        connection = connection or self.bot.pool

        query = 'SELECT role_id FROM muted_roles WHERE guild_id = $1;'
        row = await connection.fetchrow(query, guild.id)

        if not row:
            return None

        return discord.utils.get(guild.roles, id=row['role_id'])

    async def _get_muted_role(self, guild, connection=None):
        role = await self._get_muted_role_from_db(guild, connection=connection)
        if role:
            return role

        def probably_mute_role(r):
            lower = r.name.lower()
            return lower == 'muted' or 'mute' in lower

        return discord.utils.find(probably_mute_role, reversed(guild.roles))

    async def _update_muted_role(self, guild, new_role, connection=None):
        connection = connection or self.bot.pool
        query = """
            INSERT INTO   muted_roles (guild_id, role_id)
            VALUES        ($1, $2)
            ON CONFLICT   (guild_id)
            DO UPDATE SET role_id = $2;
        """
        await connection.execute(query, guild.id, new_role.id)

    @staticmethod
    async def _regen_muted_role_perms(role, *channels):
        muted_perms = dict.fromkeys([
            'send_messages', 'manage_messages', 'add_reactions', 'speak',
            'connect', 'use_voice_activation'
        ], False)

        permissions_in = channels[0].guild.me.permissions_in
        for channel in channels:
            if not permissions_in(channel).manage_roles:
                continue

            await asyncio.sleep(random.uniform(0, 0.5))

            try:
                await channel.set_permissions(role, **muted_perms)
            except discord.NotFound as e:
                if 'Unknown Overwrite' in str(e):
                    raise
            except discord.HTTPException:
                pass

    async def _do_mute(self,
                       member,
                       when,
                       role,
                       *,
                       connection=None,
                       reason=None):
        if role in member.roles:
            raise AlreadyMuted(f'{member.mention} has already been muted.')

        await member.add_roles(role, reason=reason)

        if when is not None:
            args = (member.guild.id, member.id, role.id)
            await self.bot.db_scheduler.add_abs(when, 'mute_complete', args)

    async def _create_muted_role(self, ctx):
        # Creating roles can take sooooo f*****g much time. Better release the pool.
        await ctx.release()

        bucket = self._mute_role_create_cooldowns.get_bucket(ctx.message)
        if not bucket.get_tokens():
            retry_after = bucket.update_rate_limit() or 0
            raise commands.CommandOnCooldown(bucket, retry_after)

        if not await ctx.ask_confirmation(
                'No `muted` role found. Create a new one?',
                delete_after=False):
            await ctx.send(
                f'A `muted` role couldn\'t be found. Set one with `{ctx.clean_prefix}setmuterole Role`'
            )
            return None

        bucket.update_rate_limit()
        async with ctx.typing():
            ctx.__new_mute_role_message__ = await ctx.send(
                'Creating `muted` role. Please wait...')
            role = await ctx.guild.create_role(
                name='Vale.py-Muted',
                color=discord.Color.red(),
                reason='Creating new muted role')

            with contextlib.suppress(discord.HTTPException):
                await role.edit(position=ctx.me.top_role.position - 1)

            await self._regen_muted_role_perms(role, *ctx.guild.channels)
            await ctx.acquire()
            await self._update_muted_role(ctx.guild, role, ctx.db)
            return role

    @commands.command(name='mute')
    @commands.has_permissions(manage_guild=True)
    @commands.bot_has_permissions(manage_roles=True)
    async def _mute(self,
                    ctx,
                    member: CheckedMember,
                    duration: typing.Optional[time.Delta] = None,
                    *,
                    reason: Reason = None):
        """Mutes a user for an optional amount of time."""

        reason = reason or f'By {ctx.author}'

        async def try_edit(content):
            try:
                await ctx.__new_mute_role_message__.edit(content=content)
            except (AttributeError, discord.NotFound):
                await ctx.send(content)

        role = await self._get_muted_role(ctx.guild, ctx.db)
        if role is None:
            try:
                role = await self._create_muted_role(ctx)
            except discord.NotFound:
                return await ctx.send(
                    'Please don\'t delete this role while I\'m setting it up.')
            except asyncio.TimeoutError:
                return await ctx.send('Sorry. You took too long...')
            except commands.CommandOnCooldown as e:
                return await ctx.send(
                    f'You\'re deleting the `muted` role too often. Please wait {time.duration_units(e.retry_after)} before trying again, '
                    f'or set a `muted` role with `{ctx.clean_prefix}setmuterole Role`.'
                )

            if role is None:
                return

        if duration is None:
            when = None
            for_how_long = 'permanently'
        else:
            when = ctx.message.created_at + duration.delta
            for_how_long = f'for {duration}'

        await self._do_mute(member,
                            when,
                            role,
                            connection=ctx.db,
                            reason=reason)
        await try_edit(
            f'Done. {member.mention} will now be muted {for_how_long}.')

    @_mute.error
    async def _mute_error(self, ctx, error):
        if isinstance(error, AlreadyMuted):
            await ctx.send(error)
        else:
            ctx.__bypass_local_error__ = True

    @commands.command(name='mutetime')
    async def _mute_time(self, ctx, member: discord.Member = None):
        """Shows the time left for a member's mute. Defaults to yourself."""

        member = member or ctx.author

        role = await self._get_muted_role(ctx.guild, ctx.db)
        if role not in member.roles:
            return await ctx.send(f'{member} is not muted...')

        query = """
            SELECT expires
            FROM   scheduler
            WHERE  event = 'mute_complete'
            AND    args_kwargs #>> '{args,0}' = $1
            AND    args_kwargs #>> '{args,1}' = $2

            -- The below condition is in case we have this scenario:
            --  - Member was muted
            --  - Mute role was changed while the user was muted
            --  - Member was muted again with the new role

            AND    args_kwargs #>> '{args,2}' = $3
            LIMIT  1;
        """
        entry = await ctx.db.fetchrow(query, str(ctx.guild.id), str(member.id),
                                      str(role.id))
        if not entry:
            return await ctx.send(
                f'{member} has been perm-muted. Probably the role was added manually...'
            )

        when = entry['expires']
        await ctx.send(
            f'{member} has {time.human_timedelta(when)} remaining. Unmute on {when: %c}.'
        )

    async def _remove_time_entry(self,
                                 guild,
                                 member,
                                 connection=None,
                                 *,
                                 event='mute_complete'):
        connection = connection or self.bot.pool

        query = """
            SELECT   id, expires
            FROM     scheduler
            WHERE    event = $3
            AND      args_kwargs #>> '{args,0}' = $1
            AND      args_kwargs #>> '{args,1}' = $2
            ORDER BY expires
            LIMIT    1;
        """
        entry = await connection.fetchrow(query, str(guild.id), str(member.id),
                                          event)
        if not entry:
            return None

        await self.bot.db_scheduler.remove(discord.Object(entry['id']))
        return entry['expires']

    @commands.command(name='unmute')
    @commands.has_permissions(manage_guild=True)
    @commands.bot_has_permissions(manage_roles=True)
    async def unmute(self,
                     ctx,
                     member: discord.Member,
                     *,
                     reason: Reason = None):
        """Unmutes a user."""

        reason = reason or f'Unmute by {ctx.author}'

        role = await self._get_muted_role(member.guild, ctx.db)
        if role not in member.roles:
            return await ctx.send(f'{member} hasn\'t been muted.')

        await member.remove_roles(role, reason=reason)
        await self._remove_time_entry(member.guild, member, ctx.db)
        await ctx.send(f'{member.mention} is no longer muted.')

    @commands.command(name='setmuterole', aliases=['muterole'])
    @commands.has_permissions(manage_guild=True, manage_roles=True)
    async def _set_mute_role(self, ctx, *, role: discord.Role):
        """Sets the `muted` role for the server."""

        await self._update_muted_role(ctx.guild, role, ctx.db)
        await ctx.send(f'Set the `muted` role to **{role}**.')

    @commands.command(name='kick')
    @commands.has_permissions(kick_members=True)
    @commands.bot_has_permissions(kick_members=True)
    async def _kick(self,
                    ctx,
                    member: CheckedMember,
                    *,
                    reason: Reason = None):
        """Kicks a user."""

        reason = reason or f'By {ctx.author}'
        await member.kick(reason=reason)
        await ctx.message.add_reaction('\N{WHITE HEAVY CHECK MARK}')

    @commands.command(name='softban')
    @commands.has_permissions(kick_members=True, manage_guild=True)
    @commands.bot_has_permissions(ban_members=True)
    async def _soft_ban(self,
                        ctx,
                        member: CheckedMember,
                        *,
                        reason: Reason = None):
        """Softbans a user."""

        reason = reason or f'By {ctx.author}'
        await member.ban(reason=reason)
        await member.unban(reason=f'Softban (Original reason: {reason})')
        await ctx.message.add_reaction('\N{WHITE HEAVY CHECK MARK}')

    @commands.command(name='tempban')
    @commands.has_permissions(ban_members=True)
    @commands.bot_has_permissions(ban_members=True)
    async def _temp_ban(self,
                        ctx,
                        member: CheckedMember,
                        duration: time.Delta,
                        reason: Reason = None):
        """Temporarily bans a user."""

        reason = reason or f'By {ctx.author}'
        await ctx.guild.ban(member, reason=reason)
        await ctx.message.add_reaction('\N{WHITE HEAVY CHECK MARK}')

        await self.bot.db_scheduler.add(duration.delta, 'tempban_complete',
                                        (ctx.guild.id, member.id))

    @commands.command(name='ban')
    @commands.has_permissions(ban_members=True)
    @commands.bot_has_permissions(ban_members=True)
    async def _ban(self,
                   ctx,
                   member: CheckedMemberID,
                   *,
                   reason: Reason = None):
        """Bans a member.

        You can use this to ban someone even if he's not in the server, just use the ID.
        """

        reason = reason or f'By {ctx.author}'
        await ctx.guild.ban(member, reason=reason)
        await ctx.message.add_reaction('\N{WHITE HEAVY CHECK MARK}')

    @commands.command(name='unban')
    @commands.has_permissions(ban_members=True)
    @commands.bot_has_permissions(ban_members=True)
    async def _unban(self, ctx, user: BannedMember, *, reason: Reason = None):
        """Unbans the user."""

        reason = reason or f'By {ctx.author}'
        await ctx.guild.unban(user, reason=reason)
        await self._remove_time_entry(ctx.guild,
                                      user,
                                      ctx.db,
                                      event='tempban_complete')
        await ctx.message.add_reaction('\N{WHITE HEAVY CHECK MARK}')

    @commands.command(name='massban')
    @commands.has_permissions(ban_members=True)
    async def _mass_ban(self,
                        ctx,
                        members: commands.Greedy[_CheckedMember],
                        delete_days: typing.Optional[int] = 0,
                        *,
                        reason: Reason):
        """Bans multiple users from the server.
        """

        for member in members:
            await ctx.guild.ban(member,
                                reason=reason,
                                delete_message_days=delete_days)

        await ctx.message.add_reaction('\N{WHITE HEAVY CHECK MARK}')

    # Corresponding events for that crap

    @commands.Cog.listener()
    async def on_message(self, message):
        await self.check_slowmode(message)

    async def on_guild_channel_create(self, channel):
        guild = channel.guild
        role = await self._get_muted_role_from_db(guild)
        if not role:
            return

        await self._regen_muted_role_perms(role, channel)

    @commands.Cog.listener()
    async def on_member_join(self, member):
        expires = await self._remove_time_entry(member.guild, member)
        if not expires:
            return

        role = await self._get_muted_role(member.guild)
        if not role:
            return

        await self._do_mute(member,
                            expires + datetime.timedelta(seconds=3600),
                            role,
                            reason='Mute Evasion')

    @commands.Cog.listener()
    async def on_member_update(self, before, after):
        removed_roles = set(before.roles).difference(after.roles)
        if not removed_roles:
            return

        role = await self._get_muted_role(before.guild)
        if role in removed_roles:
            await self._remove_time_entry(before.guild, before)

    # And some custom events

    async def _wait_for_cache(self, name, guild_id, member_id):
        mod_log = self.bot.get_cog('ModLog')
        if mod_log:
            await mod_log.wait_for_cache(name, guild_id, member_id)

    @commands.Cog.listener()
    async def on_mute_complete(self, timer):
        guild_id, member_id, mute_role_id = timer.args
        guild = self.bot.get_guild(guild_id)
        if not guild:
            return

        member = guild.get_member(member_id)
        if not member:
            return

        role = discord.utils.get(guild.roles, id=mute_role_id)
        if not role:
            return

        await member.remove_roles(role)

    @commands.Cog.listener()
    async def on_tempban_complete(self, timer):
        guild_id, user_id = timer.args
        await self._wait_for_cache('tempban', guild_id, user_id)

        guild = self.bot.get_guild(guild_id)
        if not guild:
            return

        await guild.unban(discord.Object(user_id),
                          reason='Unban from tempban.')
Пример #40
0
class ProblemSolutionHandler(base.OperationHandler):
    SOLUTIONS_PER_PAGE = 20

    @base.require_perm(builtin.PERM_VIEW_PROBLEM_SOLUTION)
    @base.get_argument
    @base.route_argument
    @base.sanitize
    async def get(self, *, pid: document.convert_doc_id, page: int = 1):
        uid = self.user['_id'] if self.has_priv(
            builtin.PRIV_USER_PROFILE) else None
        pdoc = await problem.get(self.domain_id, pid, uid)
        if pdoc.get('hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        psdocs, pcount, pscount = await pagination.paginate(
            problem.get_multi_solution(self.domain_id, pdoc['doc_id']), page,
            self.SOLUTIONS_PER_PAGE)
        uids = {pdoc['owner_uid']}
        uids.update(psdoc['owner_uid'] for psdoc in psdocs)
        for psdoc in psdocs:
            if 'reply' in psdoc:
                uids.update(psrdoc['owner_uid'] for psrdoc in psdoc['reply'])
        udict, dudict, pssdict = await asyncio.gather(
            user.get_dict(uids),
            domain.get_dict_user_by_uid(self.domain_id, uids),
            problem.get_dict_solution_status(self.domain_id,
                                             (psdoc['doc_id']
                                              for psdoc in psdocs),
                                             self.user['_id']))
        dudict[self.user['_id']] = self.domain_user
        path_components = self.build_path(
            (self.translate('problem_main'), self.reverse_url('problem_main')),
            (pdoc['title'],
             self.reverse_url('problem_detail', pid=pdoc['doc_id'])),
            (self.translate('problem_solution'), None))
        self.render('problem_solution.html',
                    path_components=path_components,
                    pdoc=pdoc,
                    psdocs=psdocs,
                    page=page,
                    pcount=pcount,
                    pscount=pscount,
                    udict=udict,
                    dudict=dudict,
                    pssdict=pssdict)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.require_perm(builtin.PERM_CREATE_PROBLEM_SOLUTION)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_submit(self, *, pid: document.convert_doc_id, content: str):
        pdoc = await problem.get(self.domain_id, pid)
        if pdoc.get('hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        await problem.add_solution(self.domain_id, pdoc['doc_id'],
                                   self.user['_id'], content)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_edit_solution(self, *, pid: document.convert_doc_id,
                                 psid: document.convert_doc_id, content: str):
        pdoc = await problem.get(self.domain_id, pid)
        if pdoc.get('hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        psdoc = await problem.get_solution(self.domain_id, psid,
                                           pdoc['doc_id'])
        if not self.own(psdoc, builtin.PERM_EDIT_PROBLEM_SOLUTION_SELF):
            self.check_perm(builtin.PERM_EDIT_PROBLEM_SOLUTION)
        psdoc = await problem.set_solution(self.domain_id,
                                           psdoc['doc_id'],
                                           content=content)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_delete_solution(self, *, pid: document.convert_doc_id,
                                   psid: document.convert_doc_id):
        pdoc = await problem.get(self.domain_id, pid)
        if pdoc.get('hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        psdoc = await problem.get_solution(self.domain_id, psid,
                                           pdoc['doc_id'])
        if not self.own(psdoc, builtin.PERM_DELETE_PROBLEM_SOLUTION_SELF):
            self.check_perm(builtin.PERM_DELETE_PROBLEM_SOLUTION)
        await oplog.add(self.user['_id'],
                        oplog.TYPE_DELETE_DOCUMENT,
                        doc=psdoc)
        await problem.delete_solution(self.domain_id, psdoc['doc_id'])
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_edit_reply(self, *, pid: document.convert_doc_id,
                              psid: document.convert_doc_id,
                              psrid: document.convert_doc_id, content: str):
        pdoc = await problem.get(self.domain_id, pid)
        if pdoc.get('hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        psdoc, psrdoc = await problem.get_solution_reply(
            self.domain_id, psid, psrid)
        if not psdoc or psdoc['parent_doc_id'] != pdoc['doc_id']:
            raise error.DocumentNotFoundError(self.domain_id,
                                              document.TYPE_PROBLEM_SOLUTION,
                                              psid)
        if not self.own(psrdoc, builtin.PERM_EDIT_PROBLEM_SOLUTION_REPLY_SELF):
            self.check_perm(builtin.PERM_EDIT_PROBLEM_SOLUTION_REPLY)
        await problem.edit_solution_reply(self.domain_id, psid, psrid, content)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_delete_reply(self, *, pid: document.convert_doc_id,
                                psid: document.convert_doc_id,
                                psrid: document.convert_doc_id):
        pdoc = await problem.get(self.domain_id, pid)
        if pdoc.get('hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        psdoc, psrdoc = await problem.get_solution_reply(
            self.domain_id, psid, psrid)
        if not psdoc or psdoc['parent_doc_id'] != pdoc['doc_id']:
            raise error.DocumentNotFoundError(self.domain_id,
                                              document.TYPE_PROBLEM_SOLUTION,
                                              psid)
        if not self.own(psrdoc,
                        builtin.PERM_DELETE_PROBLEM_SOLUTION_REPLY_SELF):
            self.check_perm(builtin.PERM_DELETE_PROBLEM_SOLUTION_REPLY)
        await oplog.add(self.user['_id'],
                        oplog.TYPE_DELETE_SUB_DOCUMENT,
                        sub_doc=psrdoc,
                        doc_type=psdoc['doc_type'],
                        doc_id=psdoc['doc_id'])
        await problem.delete_solution_reply(self.domain_id, psid, psrid)
        self.json_or_redirect(self.url)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.require_perm(builtin.PERM_VOTE_PROBLEM_SOLUTION)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def upvote_downvote(self, *, pid: document.convert_doc_id,
                              psid: document.convert_doc_id, value: int):
        pdoc = await problem.get(self.domain_id, pid)
        if pdoc.get('hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        psdoc = await problem.get_solution(self.domain_id, psid,
                                           pdoc['doc_id'])
        psdoc, pssdoc = await problem.vote_solution(self.domain_id,
                                                    psdoc['doc_id'],
                                                    self.user['_id'], value)
        self.json_or_redirect(self.url,
                              vote=psdoc['vote'],
                              user_vote=pssdoc['vote'])

    post_upvote = functools.partialmethod(upvote_downvote, value=1)
    post_downvote = functools.partialmethod(upvote_downvote, value=-1)

    @base.require_priv(builtin.PRIV_USER_PROFILE)
    @base.require_perm(builtin.PERM_REPLY_PROBLEM_SOLUTION)
    @base.route_argument
    @base.require_csrf_token
    @base.sanitize
    async def post_reply(self, *, pid: document.convert_doc_id,
                         psid: document.convert_doc_id, content: str):
        pdoc = await problem.get(self.domain_id, pid)
        if pdoc.get('hidden', False):
            self.check_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)
        psdoc = await problem.get_solution(self.domain_id, psid,
                                           pdoc['doc_id'])
        await problem.reply_solution(self.domain_id, psdoc['doc_id'],
                                     self.user['_id'], content)
        self.json_or_redirect(self.url)
Пример #41
0
        """
        result = challenge_f()

        if result is not None:
            self.assertTrue(result)


"""
Automatically add test methods from each challenge.
"""
excluded = {"challenge", "main"}  # manually exclude some functions.
# exclude slow tests by default
slows = {
    "twentytwo", "twentyfour", "twenty",
    "thirtyeight", "thirtytwo", "thirtyone",
    "fortyseven", "fortyeight", "fiftysix"
}
for f_name, f in inspect.getmembers(matasano.challenges, inspect.isfunction):
    if not f_name.startswith("_") and f_name not in excluded:
        test_challenge = functools.partialmethod(ChallengeTestCase.challenge, challenge_f=f)
        if f_name in slows:
            test_challenge = unittest.skip("Slow test")
        setattr(
            ChallengeTestCase,
            "test_{}".format(f_name),
            test_challenge
        )

if __name__ == '__main__':
    unittest.main()
Пример #42
0
class UnifiResponse(UserList):
    ''' Wrapper around Unifi api return values '''

    def __init__(self, session, call, out):
        ''' takes the Request out and breaks it down '''

        self._client = session
        self.endpoint = call
        
        # Identifiy the correct wrapper for the return values
        # this way we can patch in helpers as needed
        data_wrapper = data_factory(call)
        try:
            self._orig = out.json()
            self.data = [ data_wrapper(session, call, x) for x in self._orig['data'] ]
        except:
            raise

        # In come cases the unifi api will return a result which does not match its 
        # count, in these cases the results have been truncated
        if 'count' in self.meta and len(self.data) != int(self.meta['count']):
            logger.warning("Truncated API response")
            self._truncated = True
        else:
            self._truncated = False

        common_keys = None
        for stuff in self.data:
            if common_keys:
                common_keys &= set(stuff.keys())
            else:
                common_keys = set(stuff.keys())
        if common_keys:
            self.keys = common_keys
        else:
            self.keys = set()

        for bar in ['key', 'name', 'mac' ]:
            if bar in self.keys:
                self.values = [ x[bar] for x in self.data ]
                break

        
    def __getitem__(self, key):
        ''' Try to act as both list and dict where possible '''
        if isinstance(key, int):
            return self.data[key]
        for keying in [ 'key', 'name', 'mac' ]:
            if keying in self.keys:
                foo = self.filter_by(keying, key, unwrap=True)
                if foo: return foo
        raise KeyError("{} not found".format(key))

    @property
    def meta(self):
        return self._orig['meta']

    @property
    def is_truncated(self):
        return self._truncated

    @property
    def is_ok(self):
        return self.meta['rc'] == 'ok'

    def filter_by(self, tag, value, unwrap=False):
        ret = list(filter(lambda x: x.get(tag,'') == value, self.data))
        if not unwrap: return ret
        if not ret: return None
        if len(ret) == 1: return ret[0]
        raise Exception("Asked to unwrap more than 1 value")

    def ifilter_by(self, tag, value, unwrap=False):
        ret = list(filter(lambda x: imatch(x.get(tag,''), value), self.data))
        if not unwrap: return ret
        if not ret: return None
        if len(ret) == 1: return ret[0]
        raise Exception("Asked to unwrap more than 1 value")

    by_name = partialmethod(filter_by, 'name')
    by_iname = partialmethod(ifilter_by, 'name')
    by_type = partialmethod(filter_by, 'type')
    by_key = partialmethod(filter_by, 'key')
Пример #43
0
def _inject_numeric_testing_methods(cls):
    def test_binop_name(suffix):
        return 'test_binop_{}_{}'.format(name, suffix)

    def test_ibinop_name(suffix):
        return 'test_ibinop_{}_{}'.format(name, suffix)

    def test_unaryop_name(suffix):
        return 'test_unaryop_{}_{}'.format(name, suffix)

    # inject testing methods for each binary operation
    for name, binop in _BINOPS:

        setattr(cls, test_binop_name('invalid_unknown'), partialmethod(_TestNumericField._test_binop_invalid_unknown, op=binop))
        setattr(cls, test_binop_name('invalid_none'), partialmethod(_TestNumericField._test_binop_invalid_none, op=binop))
        setattr(cls, test_binop_name('type_true'), partialmethod(_TestNumericField._test_binop_type_true, op=binop))
        setattr(cls, test_binop_name('type_pos_int'), partialmethod(_TestNumericField._test_binop_type_pos_int, op=binop))
        setattr(cls, test_binop_name('type_pos_vint'), partialmethod(_TestNumericField._test_binop_type_pos_vint, op=binop))
        setattr(cls, test_binop_name('value_true'), partialmethod(_TestNumericField._test_binop_value_true, op=binop))
        setattr(cls, test_binop_name('value_pos_int'), partialmethod(_TestNumericField._test_binop_value_pos_int, op=binop))
        setattr(cls, test_binop_name('value_pos_vint'), partialmethod(_TestNumericField._test_binop_value_pos_vint, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_true'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_true, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_pos_int'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_pos_int, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_pos_vint'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_pos_vint, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_true'), partialmethod(_TestNumericField._test_binop_lhs_value_same_true, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_pos_int'), partialmethod(_TestNumericField._test_binop_lhs_value_same_pos_int, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_pos_vint'), partialmethod(_TestNumericField._test_binop_lhs_value_same_pos_vint, op=binop))
        setattr(cls, test_binop_name('type_neg_int'), partialmethod(_TestNumericField._test_binop_type_neg_int, op=binop))
        setattr(cls, test_binop_name('type_neg_vint'), partialmethod(_TestNumericField._test_binop_type_neg_vint, op=binop))
        setattr(cls, test_binop_name('value_neg_int'), partialmethod(_TestNumericField._test_binop_value_neg_int, op=binop))
        setattr(cls, test_binop_name('value_neg_vint'), partialmethod(_TestNumericField._test_binop_value_neg_vint, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_neg_int'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_neg_int, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_neg_vint'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_neg_vint, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_neg_int'), partialmethod(_TestNumericField._test_binop_lhs_value_same_neg_int, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_neg_vint'), partialmethod(_TestNumericField._test_binop_lhs_value_same_neg_vint, op=binop))
        setattr(cls, test_binop_name('type_false'), partialmethod(_TestNumericField._test_binop_type_false, op=binop))
        setattr(cls, test_binop_name('type_zero_int'), partialmethod(_TestNumericField._test_binop_type_zero_int, op=binop))
        setattr(cls, test_binop_name('type_zero_vint'), partialmethod(_TestNumericField._test_binop_type_zero_vint, op=binop))
        setattr(cls, test_binop_name('value_false'), partialmethod(_TestNumericField._test_binop_value_false, op=binop))
        setattr(cls, test_binop_name('value_zero_int'), partialmethod(_TestNumericField._test_binop_value_zero_int, op=binop))
        setattr(cls, test_binop_name('value_zero_vint'), partialmethod(_TestNumericField._test_binop_value_zero_vint, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_false'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_false, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_zero_int'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_zero_int, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_zero_vint'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_zero_vint, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_false'), partialmethod(_TestNumericField._test_binop_lhs_value_same_false, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_zero_int'), partialmethod(_TestNumericField._test_binop_lhs_value_same_zero_int, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_zero_vint'), partialmethod(_TestNumericField._test_binop_lhs_value_same_zero_vint, op=binop))
        setattr(cls, test_binop_name('type_pos_float'), partialmethod(_TestNumericField._test_binop_type_pos_float, op=binop))
        setattr(cls, test_binop_name('type_neg_float'), partialmethod(_TestNumericField._test_binop_type_neg_float, op=binop))
        setattr(cls, test_binop_name('type_pos_vfloat'), partialmethod(_TestNumericField._test_binop_type_pos_vfloat, op=binop))
        setattr(cls, test_binop_name('type_neg_vfloat'), partialmethod(_TestNumericField._test_binop_type_neg_vfloat, op=binop))
        setattr(cls, test_binop_name('value_pos_float'), partialmethod(_TestNumericField._test_binop_value_pos_float, op=binop))
        setattr(cls, test_binop_name('value_neg_float'), partialmethod(_TestNumericField._test_binop_value_neg_float, op=binop))
        setattr(cls, test_binop_name('value_pos_vfloat'), partialmethod(_TestNumericField._test_binop_value_pos_vfloat, op=binop))
        setattr(cls, test_binop_name('value_neg_vfloat'), partialmethod(_TestNumericField._test_binop_value_neg_vfloat, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_pos_float'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_pos_float, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_neg_float'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_neg_float, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_pos_vfloat'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_pos_vfloat, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_neg_vfloat'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_neg_vfloat, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_pos_float'), partialmethod(_TestNumericField._test_binop_lhs_value_same_pos_float, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_neg_float'), partialmethod(_TestNumericField._test_binop_lhs_value_same_neg_float, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_pos_vfloat'), partialmethod(_TestNumericField._test_binop_lhs_value_same_pos_vfloat, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_neg_vfloat'), partialmethod(_TestNumericField._test_binop_lhs_value_same_neg_vfloat, op=binop))
        setattr(cls, test_binop_name('type_zero_float'), partialmethod(_TestNumericField._test_binop_type_zero_float, op=binop))
        setattr(cls, test_binop_name('type_zero_vfloat'), partialmethod(_TestNumericField._test_binop_type_zero_vfloat, op=binop))
        setattr(cls, test_binop_name('value_zero_float'), partialmethod(_TestNumericField._test_binop_value_zero_float, op=binop))
        setattr(cls, test_binop_name('value_zero_vfloat'), partialmethod(_TestNumericField._test_binop_value_zero_vfloat, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_zero_float'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_zero_float, op=binop))
        setattr(cls, test_binop_name('lhs_addr_same_zero_vfloat'), partialmethod(_TestNumericField._test_binop_lhs_addr_same_zero_vfloat, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_zero_float'), partialmethod(_TestNumericField._test_binop_lhs_value_same_zero_float, op=binop))
        setattr(cls, test_binop_name('lhs_value_same_zero_vfloat'), partialmethod(_TestNumericField._test_binop_lhs_value_same_zero_vfloat, op=binop))

    # inject testing methods for each unary operation
    for name, unaryop in _UNARYOPS:
        setattr(cls, test_unaryop_name('type'), partialmethod(_TestNumericField._test_unaryop_type, op=unaryop))
        setattr(cls, test_unaryop_name('value'), partialmethod(_TestNumericField._test_unaryop_value, op=unaryop))
        setattr(cls, test_unaryop_name('addr_same'), partialmethod(_TestNumericField._test_unaryop_addr_same, op=unaryop))
        setattr(cls, test_unaryop_name('value_same'), partialmethod(_TestNumericField._test_unaryop_value_same, op=unaryop))

    # inject testing methods for each inplace binary operation
    for name, ibinop in _IBINOPS:
        setattr(cls, test_ibinop_name('invalid_unknown'), partialmethod(_TestNumericField._test_ibinop_invalid_unknown, op=ibinop))
        setattr(cls, test_ibinop_name('invalid_none'), partialmethod(_TestNumericField._test_ibinop_invalid_none, op=ibinop))
        setattr(cls, test_ibinop_name('type_true'), partialmethod(_TestNumericField._test_ibinop_type_true, op=ibinop))
        setattr(cls, test_ibinop_name('value_true'), partialmethod(_TestNumericField._test_ibinop_value_true, op=ibinop))
        setattr(cls, test_ibinop_name('type_pos_int'), partialmethod(_TestNumericField._test_ibinop_type_pos_int, op=ibinop))
        setattr(cls, test_ibinop_name('type_pos_vint'), partialmethod(_TestNumericField._test_ibinop_type_pos_vint, op=ibinop))
        setattr(cls, test_ibinop_name('value_pos_int'), partialmethod(_TestNumericField._test_ibinop_value_pos_int, op=ibinop))
        setattr(cls, test_ibinop_name('value_pos_vint'), partialmethod(_TestNumericField._test_ibinop_value_pos_vint, op=ibinop))
        setattr(cls, test_ibinop_name('type_neg_int'), partialmethod(_TestNumericField._test_ibinop_type_neg_int, op=ibinop))
        setattr(cls, test_ibinop_name('type_neg_vint'), partialmethod(_TestNumericField._test_ibinop_type_neg_vint, op=ibinop))
        setattr(cls, test_ibinop_name('value_neg_int'), partialmethod(_TestNumericField._test_ibinop_value_neg_int, op=ibinop))
        setattr(cls, test_ibinop_name('value_neg_vint'), partialmethod(_TestNumericField._test_ibinop_value_neg_vint, op=ibinop))
        setattr(cls, test_ibinop_name('type_false'), partialmethod(_TestNumericField._test_ibinop_type_false, op=ibinop))
        setattr(cls, test_ibinop_name('value_false'), partialmethod(_TestNumericField._test_ibinop_value_false, op=ibinop))
        setattr(cls, test_ibinop_name('type_zero_int'), partialmethod(_TestNumericField._test_ibinop_type_zero_int, op=ibinop))
        setattr(cls, test_ibinop_name('type_zero_vint'), partialmethod(_TestNumericField._test_ibinop_type_zero_vint, op=ibinop))
        setattr(cls, test_ibinop_name('value_zero_int'), partialmethod(_TestNumericField._test_ibinop_value_zero_int, op=ibinop))
        setattr(cls, test_ibinop_name('value_zero_vint'), partialmethod(_TestNumericField._test_ibinop_value_zero_vint, op=ibinop))
        setattr(cls, test_ibinop_name('type_pos_float'), partialmethod(_TestNumericField._test_ibinop_type_pos_float, op=ibinop))
        setattr(cls, test_ibinop_name('type_neg_float'), partialmethod(_TestNumericField._test_ibinop_type_neg_float, op=ibinop))
        setattr(cls, test_ibinop_name('type_pos_vfloat'), partialmethod(_TestNumericField._test_ibinop_type_pos_vfloat, op=ibinop))
        setattr(cls, test_ibinop_name('type_neg_vfloat'), partialmethod(_TestNumericField._test_ibinop_type_neg_vfloat, op=ibinop))
        setattr(cls, test_ibinop_name('value_pos_float'), partialmethod(_TestNumericField._test_ibinop_value_pos_float, op=ibinop))
        setattr(cls, test_ibinop_name('value_neg_float'), partialmethod(_TestNumericField._test_ibinop_value_neg_float, op=ibinop))
        setattr(cls, test_ibinop_name('value_pos_vfloat'), partialmethod(_TestNumericField._test_ibinop_value_pos_vfloat, op=ibinop))
        setattr(cls, test_ibinop_name('value_neg_vfloat'), partialmethod(_TestNumericField._test_ibinop_value_neg_vfloat, op=ibinop))
        setattr(cls, test_ibinop_name('type_zero_float'), partialmethod(_TestNumericField._test_ibinop_type_zero_float, op=ibinop))
        setattr(cls, test_ibinop_name('type_zero_vfloat'), partialmethod(_TestNumericField._test_ibinop_type_zero_vfloat, op=ibinop))
        setattr(cls, test_ibinop_name('value_zero_float'), partialmethod(_TestNumericField._test_ibinop_value_zero_float, op=ibinop))
        setattr(cls, test_ibinop_name('value_zero_vfloat'), partialmethod(_TestNumericField._test_ibinop_value_zero_vfloat, op=ibinop))
Пример #44
0
        if desc.get('masks'):
            fast_result = fast_func(im=image, mask=desc['masks'][idx] if len(desc['masks']) > 1 else desc['masks'][0])
            slow_result = slow_func(im=image, mask=desc['masks'][idx] if len(desc['masks']) > 1 else desc['masks'][0])
        else:
            fast_result = fast_func(im=image)
            slow_result = slow_func(im=image)
    if desc.get('validator_takes_desc_and_flavor_args'):
        desc['validator'](self, fast_result, slow_result, desc=desc, flavor=flavor)
    else:
        desc['validator'](self, fast_result, slow_result)

for desc in TARGET_FUNC_DESCS:
    for flavor in IMAGE_FLAVORS:
        if 'accepted_dtypes' not in desc or any(numpy.issubdtype(flavor.dtype, accepted_dtype) for accepted_dtype in desc['accepted_dtypes']):
            if 'slow_func' in desc and 'fast_func' in desc and 'validator' in desc:
                 setattr(NDImageStatisticsTestCase, 'test_{}_{}'.format(desc['name'], flavor.name), functools.partialmethod(_test, desc=desc, flavor=flavor))
            # def test(self):
            #     if desc.
            benchmarks.append(functools.partial(_benchmark, desc, flavor))

if __name__ == '__main__':
    import sys
    if len(sys.argv) == 1 or sys.argv[1] not in ('test', 'benchmark', 'test_and_benchmark'):
        print("***Defaulting to running tests.  Supply benchmark or test_and_benchmark as an argument to do otherwise.***")
        unittest.main(argv=sys.argv)
    elif sys.argv[1] == 'test':
        unittest.main(argv=sys.argv[0:1] + sys.argv[2:])
    elif sys.argv[1] == 'benchmark':
        print("All results in milliseconds.  Running benchmarks...")
        for benchmark in benchmarks:
            benchmark()
Пример #45
0
            return normalize.get_degraded_scores([str(self)], as_lv, ui)[str(self)]

    def apostrophe(self, as_lv=None):
        if not as_lv: as_lv = self.lv
        apos_char = normalize.apostrophe(as_lv)
        return self.sub(r"'", apos_char)

    def dict(self):
        return {'langvar' : self.lv, 'txt' : str(self)}

    @classmethod
    def undict(cls, d):
        return cls(d['txt'], d['langvar'])

for i in ['__getitem__', '__mul__', '__rmul__', '__mod__', '__rmod__']:
    setattr(Ex, i, partialmethod(Ex.__special_obj__, i))

for i in ['__eq__', '__lt__', '__gt__', '__le__', '__ge__', '__ne__', '__contains__']:
    setattr(Ex, i, partialmethod(Ex.__special_cmp__, i))


class Df(Ex):
    def pretty(self, indent=0):
        ind = '  ' * indent
        out = ""
        out += ind + self.__class__.__name__.lower() + '\n'
        out += ind + '  ' + self.lv + '\n'
        out += ind + '  ' + self.text + '\n'
        return out

    @classmethod
Пример #46
0
    def decorate(class_):
        for method in method_names:
            setattr(class_, method, partialmethod(call_attribute_method, method))

        return class_
Пример #47
0
 def __getattr__(self, name):
     if name.strip('_') in self.errs:
         f = functools.partialmethod(self.propagate, name=name)
         return f
     else:
         raise AttributeError
Пример #48
0
from pymongo.errors import OperationFailure

import src
from controller import MSGHandler
from backend_modules import router
from src import messages as msg
from src.db import User, Code, NoObjectReturnedFromDB, \
    ConditionNotMetError, Course, Room
from src.pub_sub import MalformedMessageError
from src.wsclass import subscribe

_path = msg.join_path('panels', 'user')

MSGHandler.send_user_not_loaded_error = partialmethod(
    MSGHandler.send_error,
    'userNotLoaded',
    description='There was no loaded user when '
                'this message arrived.'
)

MSGHandler.send_room_not_loaded_error = partialmethod(
    MSGHandler.send_error,
    'roomNotLoaded',
    description='There was no loaded room when '
    'this message arrived.'
)


def _logout_and_close(self, reason):
    """Send a logout message and close the connection."""
    self.write_message({'type': 'logout'})
    self.close(1000, reason)
Пример #49
0
 def pass_args(*args, **kwargs):
     return functools.partialmethod(method, *args, **kwargs)
Пример #50
0
    def setup_initializer(self, initializer=None):
        if initializer is None:
            def initializer(self, *args, **kwargs):
                self.instance = self.subject(*args, **kwargs)

        self.factory.methods['__init__'] = partialmethod(initializer)
Пример #51
0
class IsiClient(object):
    ''' Bare bones Isilon RESTful client designed for utter simplicity '''
    def __init__(self,
                 server=None,
                 username='',
                 password='',
                 port=8080,
                 verify=True):
        self._cookiejar_path = os.path.expanduser('~/.isilon_cookiejar')
        self._s = requests.Session()
        # Set a fixed user-agent since the isisessid is only valid for the current
        # user-agent
        self._s.headers['User-Agent'] = 'simple-isi library based on requests'
        self._s.cookies = fcj()
        try:
            self._s.cookies.load(self._cookiejar_path,
                                 ignore_discard=True,
                                 ignore_expires=True)
        except BaseException:
            logger.warning("Could not load cookies from %s",
                           self._cookiejar_path)
        self.server = server
        self.username = username
        self.password = password
        self.port = port
        self._s.verify = verify
        # if we have cookies make an attempt to login
        logger.debug("Attempting to use cached credentials")
        #self._expires = time.time() + self.refresh_session()
        self._expires = -1
        self.is_ready(prompt=False, force=True)

    def is_ready(self, auto_refresh=600, prompt=True, force=False):
        expires_in = self._expires - time.time()
        # good session
        if not force and expires_in >= auto_refresh:
            return True
        # will expire in less than auto_refresh or forced
        # from a unknown state
        if force or (expires_in > 0 and expires_in < auto_refresh):
            self._expires = time.time() + self.refresh_session()
            return True
        # try to create a new session, will work if username and
        # password are stored
        try:
            self._expires = time.time() + self.create_session()
            return True
        except BaseException:
            pass
        # Must query
        if not prompt:
            return False
        try:
            self._expires = time.time() + self.auth()
            return True
        except BaseException:
            pass
        return False

    def refresh_session(self):
        # returns the time remaining in the session
        try:
            out = self.get('session/1/session', ready_check=False)
            data = out.json()
            self.username = data['username']
            length = min(data["timeout_absolute"], data["timeout_inactive"])
            if length > 60:
                # good senssion
                self._s.cookies.save(self._cookiejar_path,
                                     ignore_discard=True,
                                     ignore_expires=True)
                return length
        except BaseException:
            # Fallthough for errors
            return -1

    def create_session(self):
        # attempts to authenticate with session module
        if self.username == '' or self.password == '':
            raise ValueError("Can't login without credentials")
        login = {
            'username': self.username,
            'password': self.password,
            'services': ['platform', 'namespace']
        }
        try:
            self.post('session/1/session', json=login, ready_check=False)
            self.ready = True
        except BaseException:
            logger.debug("Login failure")
        return self.refresh_session()

    def logout(self):
        ''' logout of a valid session / destroy cookie and 
            authentication tokens '''
        self.password = ''
        self._expires = -1
        out = self.delete('session/1/session',
                          ready_check=False,
                          raise_on_error=False)

    def __repr__(self):
        if self._s.auth:
            auth = "with auth of {}/{}".format(self.username,
                                               "*" * len(self.password))
        else:
            auth = "NO AUTHENTICATION TOKEN"
        return "IsiClient-https://{}:{} {}".format(self.server, self.port,
                                                   auth)

    def auth(self):
        # Query for interactive credentials

        # only works for ttys
        if not sys.stdin.isatty():
            logger.warning(
                "Session not ready and no interactive credentials, this will probably fail"
            )
        #    return None

        # Start interactive login
        print(
            "Please enter credentials for Isilon https://{}:{}\nUsername (CR={}): "
            .format(self.server, self.port, getuser()),
            file=sys.stderr,
            end='')
        username = input()
        if username == "":
            username = getuser()
        password = getpass("{} Password : "******"Unauthenticates REST call, will probably fail")

        if x_append_prefix:
            url = 'https://{}:{}/{}'.format(self.server, self.port,
                                            quote(endpoint))
        else:
            url = endpoint
        logger.debug("%s %s <- %s", method, url, repr(json)[:20])
        out = self._s.request(method,
                              url,
                              json=json,
                              stream=stream,
                              params=params)
        # monkeypatch for iterating over the Isilon data
        out.iter_json = partial(self.iter_out, out)
        logger.debug("Results from %s status %i preview %s", out.url,
                     out.status_code, out.text[:20])
        if raise_on_error and out.status_code != requests.codes.ok:
            raise IsiApiError(out)
        return out

    # Primary calls are just wrappers around request
    get = partialmethod(request, 'GET')
    head = partialmethod(request, 'HEAD')
    post = partialmethod(request, 'POST')
    delete = partialmethod(request, 'DELETE')

    @staticmethod
    def get_resume_id(out):
        try:
            return out.json()['resume']
        except BaseException:
            return None

    @staticmethod
    def find_collection(data):
        ''' Find the collection in the json '''
        if 'resume' in data:
            for key, value in data.items():
                if isinstance(value, list):
                    return key
        if 'directory' in data:
            return 'directory'
        if 'children' in data:
            return 'children'
        if 'summary' in data:
            return 'summary'
        return None

    def iter_out(self, out, tag=None):
        ''' Page through the results on the named tag
            attempt to guess the f*g if not give '''
        data = out.json()

        if not tag:
            # autodetect tag
            tag = IsiClient.find_collection(data)

        if not tag:
            yield data
            return

        # Page through results yielding the tag
        # this should be yield from but that's not py2 safe
        for page in self.page_out(out):
            # python3 only, so unwind it
            # yield from page.json()[tag]
            for item in page.json()[tag]:
                yield item

    def page_out(self, out):
        ''' helper to page results that have a resume entity '''

        # Return the input
        yield out
        resume = IsiClient.get_resume_id(out)
        try:
            url = out.url.split('?', 1)[0]
        except BaseException:
            url = out.url
        # While the input containes a resume token keep refreshing it
        # via additional get calls
        while resume:
            out = self.get(url, x_append_prefix=False, resume=resume)
            yield out
            resume = IsiClient.get_resume_id(out)
Пример #52
0
 def __new__(cls, base, *args, **kwargs):
     return type(
         base.__name__, (base, ),
         {'__init__': partialmethod(base.__init__, *args, **kwargs)})
Пример #53
0
class ApiClient(SyncContextManager):
    log = LOG.getChild('ApiClient')
    _max_conns = 20
    _health_check_interval = 30

    def __init__(self, env):
        self.env = env
        self.session = None
        self._tmpdir = None
        self._sock = None
        self._proxy = None
        self._closed = False
        self._startup_lock = curio.Lock()
        self._poll_lock = curio.Lock()
        self._apiVersion_cache = {}
        self._kind_cache = {}
        self.healthy = curio.Event()

    async def _poll_proxy(self, wait=False):
        async with self._poll_lock:
            if not self._proxy:
                return False
            p = self._proxy
            try:
                if wait:
                    await p.wait()
                else:
                    p.poll()
            except ProcessLookupError:
                self.log.warning(f'subprocess with pid {p.pid} not found')
            else:
                if p.returncode is None:
                    return True
                if p.returncode == 0:
                    self.log.debug(f'subprocess (pid {p.pid}) exited cleanly')
                else:
                    lvl = logging.DEBUG if self._closed else logging.WARNING
                    self.log.log(
                        lvl, f'subprocess (pid {p.pid}) terminated '
                        f'with return code {p.returncode}')
            self._proxy = None
            return False

    async def _read_proxy_stdout(self):
        if not self._proxy:
            return
        try:
            async for _ in self._proxy.stdout:
                pass
        except Exception:
            pass
        finally:
            await self._poll_proxy()

    async def _wait_for_proxy(self):
        assert self._startup_lock.locked()
        line = await self._proxy.stdout.readline()
        if not line.startswith(b'Starting to serve'):
            raise ApiError(reason='Got gibberish from kubectl proxy')
        await curio.spawn(self._read_proxy_stdout, daemon=True)

    async def _ensure_proxy(self):
        if self._closed:
            raise RuntimeError('API client is closed')
        assert self._startup_lock.locked()
        if await self._poll_proxy():
            return

        if not self._tmpdir:
            self._tmpdir = tempfile.TemporaryDirectory()
            self._sock = os.path.join(self._tmpdir.name, 'proxy.sock')

        self._proxy = self.env.spawn_kubectl(
            ['proxy', '-u', self._sock],
            stdout=subprocess.PIPE,
            preexec_fn=os.setpgrp,
        )

        await self._wait_for_proxy()

    async def _ensure_session(self):
        async with self._startup_lock:
            await self._ensure_proxy()
            if self.session is None:
                self.session = UnixSession(self._sock,
                                           connections=self._max_conns)

    def _cleanup_tmpdir(self):
        self._sock = None
        if self._tmpdir:
            self._tmpdir.cleanup()
            self._tmpdir = None

    def __del__(self):
        if self._proxy:
            try:
                self._proxy.terminate()
            except ProcessLookupError:
                pass
        self._closed = True
        self._cleanup_tmpdir()

    @sync_wrap
    async def close(self):
        self.__del__()
        if self._proxy:
            await self._poll_proxy(wait=True)

    async def __aenter__(self):
        return self

    async def __aexit__(self, typ, val, tb):
        await self.close()

    async def _kindify(self, obj):
        if 'apiVersion' in obj and 'kind' in obj:
            try:
                kindCls = await self.getKind(obj['apiVersion'], obj['kind'])
            except (ApiVersionNotFoundError, ResourceKindNotFoundError):
                pass
            else:
                obj = kindCls(obj)
        return obj

    @sync_wrap
    async def request(self, method, path, **kw):
        self.log.debug(f'{method} {path}')
        await self._ensure_session()
        resp = await self.session.request(method=method, path=path, **kw)
        if not (200 <= resp.status_code < 300):
            if resp.status_code == 404:
                raise ResourceNotFoundError(resp)
            if resp.headers.get('content-type',
                                '').split(';', 1)[0] == 'application/json':
                j = resp.json()
                reason = j.get('message')
            else:
                reason = None
            raise ApiError(resp, reason)
        if kw.get('stream'):
            return StreamIter(self, resp)
        if resp.headers.get('content-type',
                            '').split(';', 1)[0] == 'application/json':
            return await self._kindify(resp.json())
        return resp.text

    get = partialmethod(request, 'GET')
    head = partialmethod(request, 'HEAD')
    post = partialmethod(request, 'POST')
    put = partialmethod(request, 'PUT')
    patch = partialmethod(request, 'PATCH')
    delete = partialmethod(request, 'DELETE')
    options = partialmethod(request, 'OPTIONS')

    async def check_health(self):
        try:
            healthy = (await self.get('/healthz') == 'ok')
        except ApiError:
            healthy = False
        if healthy:
            await self.healthy.set()
        else:
            self.healthy.clear()
        return healthy

    async def wait_until_healthy(self):
        while not await self.check_health():
            await curio.sleep(self._health_check_interval)

    async def _get_apiVersion(self, apiVersion):
        path = f"/api{'' if apiVersion == 'v1' else 's'}/{apiVersion}"
        if path not in self._apiVersion_cache:
            try:
                res = APIResourceList(await self.get(path))
            except ApiError as err:
                if err.resp.status_code == 404:
                    raise ApiVersionNotFoundError(
                        resp=err.resp, reason='apiVersion not found') from None
                raise
            res._validate()
            self._apiVersion_cache[path] = res
            self._kind_cache.update(
                dict(((apiVersion, r['kind']), r) for r in res.resources
                     if '/' not in r['name']))
        return self._apiVersion_cache[path]

    async def _get_resourceKind(self, apiVersion, kind):
        await self._get_apiVersion(apiVersion)
        try:
            return self._kind_cache[apiVersion, kind]
        except KeyError:
            raise ResourceKindNotFoundError(
                reason='resource kind not found') from None

    @sync_wrap
    async def getKind(self, apiVersion, kind):
        try:
            return Kind.getKind(apiVersion, kind)
        except KeyError:
            res = await self._get_resourceKind(apiVersion, kind)
            return Kind.from_apiresource(apiVersion, res)

    # LEGACY STUFF - need to update (and add tests)

    async def _last_applied(self, obj):
        try:
            obj = await self._kindify(
                json.loads(
                    obj['metadata']['annotations']
                    ['kubectl.kubernetes.io/last-applied-configuration']))
        except KeyError:
            return obj
        if not obj['metadata']['annotations']:
            del obj['metadata']['annotations']
        if not obj._namespaced:
            if 'namespace' in obj['metadata']:
                del obj['metadata']['namespace']
        return obj

    @sync_wrap
    async def list(self,
                   apiVersion=None,
                   kind=None,
                   namespace=None,
                   name=None,
                   last_applied=False,
                   **kwargs):
        if apiVersion is None:
            versions = (await self.get('/api'))['versions']
            versions.extend(v['groupVersion']
                            for g in (await self.get('/apis'))['groups']
                            for v in g['versions'])
        else:
            versions = [apiVersion]
        for apiVersion in versions:
            if kind is None:
                await self._get_apiVersion(apiVersion)
                kinds = sorted(k for v, k in self._kind_cache
                               if v == apiVersion)
            else:
                kinds = [kind]
            for k in kinds:
                kindCls = await self.getKind(apiVersion, k)
                if namespace and not kindCls._namespaced:
                    continue
                path = kindCls._makeLink(name, namespace, **kwargs)
                try:
                    res = await self.get(path)
                except ApiError as e:
                    if e.resp.status_code in (404, 405):
                        continue
                    raise
                if res.get('kind', '').endswith('List'):
                    for obj in res.get('items', []):
                        obj['apiVersion'] = apiVersion
                        obj['kind'] = kind
                        if last_applied:
                            obj = await self._last_applied(obj)
                        else:
                            obj = await self._kindify(obj)
                        yield obj
                else:
                    if last_applied:
                        res = await self._last_applied(res)
                    yield res

    @sync_wrap
    async def ref_to_path(self, ref):
        if not isinstance(ref, Kind):
            kindCls = await self.getKind(ref['apiVersion'], ref['kind'])
            ref = kindCls(ref)
        return ref._selfLink()

    @sync_wrap
    async def get_refs(self, refs, last_applied=False):
        for ref in flatten_kube_lists(refs):
            if not isinstance(ref, Kind):
                try:
                    kindCls = await self.getKind(ref['apiVersion'],
                                                 ref['kind'])
                except (ApiVersionNotFoundError, ResourceKindNotFoundError):
                    continue
                ref = kindCls(ref)
            path = ref._selfLink()
            try:
                obj = await self.get(path)
            except ResourceNotFoundError:
                continue
            if last_applied:
                obj = await self._last_applied(obj)
            yield obj

    @sync_wrap
    async def create(self, obj):
        if not isinstance(obj, Kind):
            obj = await self._kindify(obj)
        kw = {}
        if 'namespace' in obj.metadata:
            kw['namespace'] = obj.metadata.namespace
        path = obj._makeLink(**kw)
        return await self.post(path, json=obj)

    @sync_wrap
    async def ensure_object(self, obj):
        try:
            await self.create(obj)
        except ApiError as err:
            if err.resp.status_code != 409:
                raise

    @sync_wrap
    async def ensure_objects(self, objs):
        for obj in objs:
            await self.ensure_object(obj)

    @sync_wrap
    async def replace(self, obj):
        if not isinstance(obj, Kind):
            obj = await self._kindify(obj)
        path = obj._selfLink()
        return await self.put(path, json=obj)

    @sync_wrap
    async def build_path(self,
                         apiVersion,
                         kind=None,
                         namespace=None,
                         name=None,
                         verb=None,
                         **kwargs):
        if kind:
            k = await self.getKind(apiVersion, kind)
            return k._makeLink(namespace=namespace,
                               name=name,
                               verb=verb,
                               **kwargs)
        return f"/api{'' if apiVersion == 'v1' else 's'}/{apiVersion}"
Пример #54
0
    'uncon_vars',
    'tagsin_mask',
    'tagsout_mask',
    'tagsin_lhs_vars',
    'tagsout_lhs_vars',
    'should_filter',
    'filter_needs_preds',
    'functionally_determines',
    'get_priority',
    'get_code',
    'rename_lhs_vars',
    'rename_rhs_rel',
    'singletonize',
    'subtract',
    ]:
    setattr(ClauseVisitor, op, partialmethod(ClauseVisitor.visit, op))


def assert_unique(vars):
    """Given a tuple of LHS vars, raise AssertionError if there are
    repeated vars (representing an equality pattern).
    """
    counts = Counter(vars)
    if not all(c == 1 for c in counts.values()):
        dups = [x for x, c in counts.items() if c > 1]
        raise AssertionError('LHS vars appear multiple times in same '
                             'clause: ' + ', '.join(dups))


class RelMemberHandler(ClauseHandler):
    
Пример #55
0
class RedisCache:
    """
    A simplified interface for a Redis connection.

    We implement several convenient methods that are fairly similar to have a dict
    behaves, and should be familiar to Python users. The biggest difference is that
    all the public methods in this class are coroutines, and must be awaited.

    Because of limitations in Redis, this cache will only accept strings, integers and
    floats both for keys and values.

    Please note that this class MUST be created as a class attribute, and that that class
    must also contain an attribute with an instance of our Bot. See `__get__` and `__set_name__`
    for more information about how this works.

    Simple example for how to use this:

    class SomeCog(Cog):
        # To initialize a valid RedisCache, just add it as a class attribute here.
        # Do not add it to the __init__ method or anywhere else, it MUST be a class
        # attribute. Do not pass any parameters.
        cache = RedisCache()

        async def my_method(self):

            # Now we're ready to use the RedisCache.
            # One thing to note here is that this will not work unless
            # we access self.cache through an _instance_ of this class.
            #
            # For example, attempting to use SomeCog.cache will _not_ work,
            # you _must_ instantiate the class first and use that instance.
            #
            # Now we can store some stuff in the cache just by doing this.
            # This data will persist through restarts!
            await self.cache.set("key", "value")

            # To get the data, simply do this.
            value = await self.cache.get("key")

            # Other methods work more or less like a dictionary.
            # Checking if something is in the cache
            await self.cache.contains("key")

            # iterating the cache
            async for key, value in self.cache.items():
                print(value)

            # We can even iterate in a comprehension!
            consumed = [value async for key, value in self.cache.items()]
    """

    _namespaces = []

    def __init__(self) -> None:
        """Initialize the RedisCache."""
        self._namespace = None
        self.bot = None
        self._increment_lock = None

    def _set_namespace(self, namespace: str) -> None:
        """Try to set the namespace, but do not permit collisions."""
        log.trace(f"RedisCache setting namespace to {namespace}")
        self._namespaces.append(namespace)
        self._namespace = namespace

    @staticmethod
    def _to_typestring(key_or_value: RedisKeyOrValue,
                       prefixes: _PrefixTuple) -> str:
        """Turn a valid Redis type into a typestring."""
        for prefix, _type in prefixes:
            if isinstance(key_or_value, _type):
                return f"{prefix}{key_or_value}"
        raise TypeError(
            f"RedisCache._to_typestring only supports the following: {prefixes}."
        )

    @staticmethod
    def _from_typestring(key_or_value: Union[bytes, str],
                         prefixes: _PrefixTuple) -> RedisKeyOrValue:
        """Deserialize a typestring into a valid Redis type."""
        # Stuff that comes out of Redis will be bytestrings, so let's decode those.
        if isinstance(key_or_value, bytes):
            key_or_value = key_or_value.decode('utf-8')

        # Now we convert our unicode string back into the type it originally was.
        for prefix, _type in prefixes:
            if key_or_value.startswith(prefix):
                return _type(key_or_value[len(prefix):])
        raise TypeError(
            f"RedisCache._from_typestring only supports the following: {prefixes}."
        )

    # Add some nice partials to call our generic typestring converters.
    # These are basically methods that will fill in some of the parameters for you, so that
    # any call to _key_to_typestring will be like calling _to_typestring with the two parameters
    # at `prefixes` and `types_string` pre-filled.
    #
    # See https://docs.python.org/3/library/functools.html#functools.partialmethod
    _key_to_typestring = partialmethod(_to_typestring, prefixes=_KEY_PREFIXES)
    _value_to_typestring = partialmethod(_to_typestring,
                                         prefixes=_VALUE_PREFIXES)
    _key_from_typestring = partialmethod(_from_typestring,
                                         prefixes=_KEY_PREFIXES)
    _value_from_typestring = partialmethod(_from_typestring,
                                           prefixes=_VALUE_PREFIXES)

    def _dict_from_typestring(self, dictionary: Dict) -> Dict:
        """Turns all contents of a dict into valid Redis types."""
        return {
            self._key_from_typestring(key): self._value_from_typestring(value)
            for key, value in dictionary.items()
        }

    def _dict_to_typestring(self, dictionary: Dict) -> Dict:
        """Turns all contents of a dict into typestrings."""
        return {
            self._key_to_typestring(key): self._value_to_typestring(value)
            for key, value in dictionary.items()
        }

    async def _validate_cache(self) -> None:
        """Validate that the RedisCache is ready to be used."""
        if self._namespace is None:
            error_message = (
                "Critical error: RedisCache has no namespace. "
                "This object must be initialized as a class attribute.")
            log.error(error_message)
            raise NoNamespaceError(error_message)

        if self.bot is None:
            error_message = (
                "Critical error: RedisCache has no `Bot` instance. "
                "This happens when the class RedisCache was created in doesn't "
                "have a Bot instance. Please make sure that you're instantiating "
                "the RedisCache inside a class that has a Bot instance attribute."
            )
            log.error(error_message)
            raise NoBotInstanceError(error_message)

        if not self.bot.redis_closed:
            await self.bot.redis_ready.wait()

    def __set_name__(self, owner: Any, attribute_name: str) -> None:
        """
        Set the namespace to Class.attribute_name.

        Called automatically when this class is constructed inside a class as an attribute.

        This class MUST be created as a class attribute in a class, otherwise it will raise
        exceptions whenever a method is used. This is because it uses this method to create
        a namespace like `MyCog.my_class_attribute` which is used as a hash name when we store
        stuff in Redis, to prevent collisions.
        """
        self._set_namespace(f"{owner.__name__}.{attribute_name}")

    def __get__(self, instance: RedisCache, owner: Any) -> RedisCache:
        """
        This is called if the RedisCache is a class attribute, and is accessed.

        The class this object is instantiated in must contain an attribute with an
        instance of Bot. This is because Bot contains our redis_session, which is
        the mechanism by which we will communicate with the Redis server.

        Any attempt to use RedisCache in a class that does not have a Bot instance
        will fail. It is mostly intended to be used inside of a Cog, although theoretically
        it should work in any class that has a Bot instance.
        """
        if self.bot:
            return self

        if self._namespace is None:
            error_message = "RedisCache must be a class attribute."
            log.error(error_message)
            raise NoNamespaceError(error_message)

        if instance is None:
            error_message = (
                "You must access the RedisCache instance through the cog instance "
                "before accessing it using the cog's class object.")
            log.error(error_message)
            raise NoParentInstanceError(error_message)

        for attribute in vars(instance).values():
            if isinstance(attribute, Bot):
                self.bot = attribute
                self._redis = self.bot.redis_session
                return self
        else:
            error_message = (
                "Critical error: RedisCache has no `Bot` instance. "
                "This happens when the class RedisCache was created in doesn't "
                "have a Bot instance. Please make sure that you're instantiating "
                "the RedisCache inside a class that has a Bot instance attribute."
            )
            log.error(error_message)
            raise NoBotInstanceError(error_message)

    def __repr__(self) -> str:
        """Return a beautiful representation of this object instance."""
        return f"RedisCache(namespace={self._namespace!r})"

    async def set(self, key: RedisKeyType, value: RedisValueType) -> None:
        """Store an item in the Redis cache."""
        await self._validate_cache()

        # Convert to a typestring and then set it
        key = self._key_to_typestring(key)
        value = self._value_to_typestring(value)

        log.trace(f"Setting {key} to {value}.")
        await self._redis.hset(self._namespace, key, value)

    async def get(self,
                  key: RedisKeyType,
                  default: Optional[RedisValueType] = None
                  ) -> Optional[RedisValueType]:
        """Get an item from the Redis cache."""
        await self._validate_cache()
        key = self._key_to_typestring(key)

        log.trace(f"Attempting to retrieve {key}.")
        value = await self._redis.hget(self._namespace, key)

        if value is None:
            log.trace(f"Value not found, returning default value {default}")
            return default
        else:
            value = self._value_from_typestring(value)
            log.trace(f"Value found, returning value {value}")
            return value

    async def delete(self, key: RedisKeyType) -> None:
        """
        Delete an item from the Redis cache.

        If we try to delete a key that does not exist, it will simply be ignored.

        See https://redis.io/commands/hdel for more info on how this works.
        """
        await self._validate_cache()
        key = self._key_to_typestring(key)

        log.trace(f"Attempting to delete {key}.")
        return await self._redis.hdel(self._namespace, key)

    async def contains(self, key: RedisKeyType) -> bool:
        """
        Check if a key exists in the Redis cache.

        Return True if the key exists, otherwise False.
        """
        await self._validate_cache()
        key = self._key_to_typestring(key)
        exists = await self._redis.hexists(self._namespace, key)

        log.trace(
            f"Testing if {key} exists in the RedisCache - Result is {exists}")
        return exists

    async def items(self) -> ItemsView:
        """
        Fetch all the key/value pairs in the cache.

        Returns a normal ItemsView, like you would get from dict.items().

        Keep in mind that these items are just a _copy_ of the data in the
        RedisCache - any changes you make to them will not be reflected
        into the RedisCache itself. If you want to change these, you need
        to make a .set call.

        Example:
        items = await my_cache.items()
        for key, value in items:
            # Iterate like a normal dictionary
        """
        await self._validate_cache()
        items = self._dict_from_typestring(await
                                           self._redis.hgetall(self._namespace
                                                               )).items()

        log.trace(
            f"Retrieving all key/value pairs from cache, total of {len(items)} items."
        )
        return items

    async def length(self) -> int:
        """Return the number of items in the Redis cache."""
        await self._validate_cache()
        number_of_items = await self._redis.hlen(self._namespace)
        log.trace(f"Returning length. Result is {number_of_items}.")
        return number_of_items

    async def to_dict(self) -> Dict:
        """Convert to dict and return."""
        return {key: value for key, value in await self.items()}

    async def clear(self) -> None:
        """Deletes the entire hash from the Redis cache."""
        await self._validate_cache()
        log.trace("Clearing the cache of all key/value pairs.")
        await self._redis.delete(self._namespace)

    async def pop(self,
                  key: RedisKeyType,
                  default: Optional[RedisValueType] = None) -> RedisValueType:
        """Get the item, remove it from the cache, and provide a default if not found."""
        log.trace(f"Attempting to pop {key}.")
        value = await self.get(key, default)

        log.trace(
            f"Attempting to delete item with key '{key}' from the cache. "
            "If this key doesn't exist, nothing will happen.")
        await self.delete(key)

        return value

    async def update(self, items: Dict[RedisKeyType, RedisValueType]) -> None:
        """
        Update the Redis cache with multiple values.

        This works exactly like dict.update from a normal dictionary. You pass
        a dictionary with one or more key/value pairs into this method. If the keys
        do not exist in the RedisCache, they are created. If they do exist, the values
        are updated with the new ones from `items`.

        Please note that keys and the values in the `items` dictionary
        must consist of valid RedisKeyTypes and RedisValueTypes.
        """
        await self._validate_cache()
        log.trace(f"Updating the cache with the following items:\n{items}")
        await self._redis.hmset_dict(self._namespace,
                                     self._dict_to_typestring(items))

    async def increment(self,
                        key: RedisKeyType,
                        amount: Optional[int, float] = 1) -> None:
        """
        Increment the value by `amount`.

        This works for both floats and ints, but will raise a TypeError
        if you try to do it for any other type of value.

        This also supports negative amounts, although it would provide better
        readability to use .decrement() for that.
        """
        log.trace(
            f"Attempting to increment/decrement the value with the key {key} by {amount}."
        )

        # We initialize the lock here, because we need to ensure we get it
        # running on the same loop as the calling coroutine.
        #
        # If we initialized the lock in the __init__, the loop that the coroutine this method
        # would be called from might not exist yet, and so the lock would be on a different
        # loop, which would raise RuntimeErrors.
        if self._increment_lock is None:
            self._increment_lock = asyncio.Lock()

        # Since this has several API calls, we need a lock to prevent race conditions
        async with self._increment_lock:
            value = await self.get(key)

            # Can't increment a non-existing value
            if value is None:
                error_message = "The provided key does not exist!"
                log.error(error_message)
                raise KeyError(error_message)

            # If it does exist, and it's an int or a float, increment and set it.
            if isinstance(value, int) or isinstance(value, float):
                value += amount
                await self.set(key, value)
            else:
                error_message = "You may only increment or decrement values that are integers or floats."
                log.error(error_message)
                raise TypeError(error_message)

    async def decrement(self,
                        key: RedisKeyType,
                        amount: Optional[int, float] = 1) -> None:
        """
        Decrement the value by `amount`.

        Basically just does the opposite of .increment.
        """
        await self.increment(key, -amount)
Пример #56
0
class GraphNode:
    """Node in the SWH graph"""
    def __init__(self, graph, node_id):
        self.graph = graph
        self.id = node_id

    def children(self):
        return Neighbors(
            self.graph,
            self.graph.java_graph.successors(self.id),
            lambda: self.graph.java_graph.outdegree(self.id),
        )

    def parents(self):
        return Neighbors(
            self.graph,
            self.graph.java_graph.predecessors(self.id),
            lambda: self.graph.java_graph.indegree(self.id),
        )

    def simple_traversal(self, ttype, direction="forward", edges="*"):
        for node in call_async_gen(self.graph.backend.simple_traversal, ttype,
                                   direction, edges, self.id):
            yield self.graph[node]

    def leaves(self, *args, **kwargs):
        yield from self.simple_traversal("leaves", *args, **kwargs)

    def visit_nodes(self, *args, **kwargs):
        yield from self.simple_traversal("visit_nodes", *args, **kwargs)

    def visit_edges(self, direction="forward", edges="*"):
        for src, dst in call_async_gen(self.graph.backend.visit_edges,
                                       direction, edges, self.id):
            yield (self.graph[src], self.graph[dst])

    def visit_paths(self, direction="forward", edges="*"):
        for path in call_async_gen(self.graph.backend.visit_paths, direction,
                                   edges, self.id):
            yield [self.graph[node] for node in path]

    def walk(self, dst, direction="forward", edges="*", traversal="dfs"):
        for node in call_async_gen(self.graph.backend.walk, direction, edges,
                                   traversal, self.id, dst):
            yield self.graph[node]

    def _count(self, ttype, direction="forward", edges="*"):
        return self.graph.backend.count(ttype, direction, edges, self.id)

    count_leaves = functools.partialmethod(_count, ttype="leaves")
    count_neighbors = functools.partialmethod(_count, ttype="neighbors")
    count_visit_nodes = functools.partialmethod(_count, ttype="visit_nodes")

    @property
    def pid(self):
        return self.graph.node2pid[self.id]

    @property
    def kind(self):
        return self.pid.split(":")[2]

    def __str__(self):
        return self.pid

    def __repr__(self):
        return "<{}>".format(self.pid)

    def dot_fragment(self):
        swh, version, kind, hash = self.pid.split(":")
        label = "{}:{}..{}".format(kind, hash[0:2], hash[-2:])
        url = BASE_URL + KIND_TO_URL_FRAGMENT[kind].format(hash)
        shape = KIND_TO_SHAPE[kind]
        return '{} [label="{}", href="{}", target="_blank", shape="{}"];'.format(
            self.id, label, url, shape)

    def _repr_svg_(self):
        nodes = [self, *list(self.children()), *list(self.parents())]
        dot = graph_dot(nodes)
        svg = dot_to_svg(dot)
        return svg
Пример #57
0
class NsClient(object):
    ''' Basic namespace client takes an IsiClient and the namespace
        prefix as setup '''
    def __init__(self, client, prefix):
        self.client = client
        self.prefix = prefix

    def call(self, method, path, **params):
        # call the papi
        return self.client.request(method,
                                   '/'.join(['namespace', self.prefix,
                                             path]), **params)

    get = partialmethod(call, 'GET')
    head = partialmethod(call, 'HEAD')
    post = partialmethod(call, 'POST')
    delete = partialmethod(call, 'DELETE')

    def scandir(self,
                path,
                detail=[
                    'name', 'owner', 'group', 'type', 'mode', 'size',
                    'block_size', 'mtime_val', 'atime_val', 'ctime_val',
                    'btime_val', 'uid', 'gid', 'id', 'nlink'
                ]):
        ''' iterate over a directory.
            workalike to os.scandir '''
        out = self.get(path, detail=','.join(detail))
        if out.headers['x-isi-ifs-target-type'] != 'container':
            raise ValueError("NOT DIRECTORY: {}".format(path))
        # py2 compatibility since no yield from
        for item in out.iter_json():
            yield item

    def walk(self, top, topdown=True, **scandir_options):
        ''' walk the directory tree
            workalike to os.walk '''
        dir_stack = [top]
        while dir_stack:
            cur = dir_stack.pop(0)
            dirs = []
            files = []
            for item in self.scandir(cur, **scandir_options):
                if 'type' in item and item['type'] == 'container':
                    dirs.append(item)
                else:
                    files.append(item)
            yield cur, dirs, files
            for item in dirs:
                dir_stack.append(os.path.join(cur, item['name']))

    @staticmethod
    def expand_dirent(entry):
        ''' Utility function to interpret the json values as python where approprate '''

        # these should be int() converted
        convert_to_int = set(
            ('size', 'block_size', 'mtime_val', 'atime_val', 'ctime_val',
             'btime_val', 'uid', 'gid', 'id', 'nlink'))
        # there should be fromtimestamp() converted
        convert_to_time = set(
            ('mtime_val', 'atime_val', 'ctime_val', 'btime_val'))
        # this is the mapping of type to mode
        type_to_flag = {
            'container': 0o0040000,
            'object': 0o0100000,
            'pipe': 0o0010000,
            'character_device': 0o0020000,
            'block_device': 0o0060000,
            'symbolic_link': 0o0120000,
            'socket': 0o0140000,
            'whiteout_file': 0
        }
        for key in convert_to_int:
            if key in entry:
                entry[key] = int(entry[key])
        for key in convert_to_time:
            if key in entry:
                entry[key] = datetime.fromtimestamp(entry[key])
        if 'size' in entry:
            entry['hsize'] = sfmt(entry['size'])
        if 'mode' in entry:
            entry['mode'] = int(entry['mode'], 8)
            try:
                entry['mode'] = entry['mode'] | type_to_flag[entry['type']]
            except BaseException:
                pass
            entry['mode_str'] = filemode(entry['mode'])

    def ll(self, path):
        ''' list a single directory '''
        for entry in self.scandir(path):
            NsClient.expand_dirent(entry)
            print(lsfmt.format(**entry))

    def llr(self, top):
        ''' list a directory recursively '''
        for path, dirs, files in self.walk(top):
            print("DIR: ", path)
            for entry in dirs:
                NsClient.expand_dirent(entry)
                print(lsfmt.format(**entry))
            for entry in files:
                NsClient.expand_dirent(entry)
                print(lsfmt.format(**entry))
Пример #58
0
    return getattr(self.value,name)(other)

def mixed_scalar_binary_method_inplace(self, other, name2):
    result = getattr(self.value,name2)(other)
    if result is NotImplemented:
        return NotImplemented
    self._monitor.set_path(self._path, result)
    return result

class MixedScalar(MixedBase):
    def __getitem__(self, item):
        value = self.value
        if not isinstance(value, (str, np.ndarray)):
            raise TypeError(type(value))
        return value[item]

for name in binary_special_method_names:
    if name.startswith("__i"):
        name2 = "__" + name[3:]
        m = partialmethod(mixed_scalar_binary_method_inplace, name2=name2)
    else:
        m = partialmethod(mixed_scalar_binary_method, name=name)
    setattr(MixedScalar, name, m)


from .MixedDict import MixedDict
from .MixedObject import MixedObject
from .Monitor import Monitor
from .Backend import Backend, DefaultBackend, CellBackend
from .get_form import is_contiguous, is_unsigned
Пример #59
0
 def __new__(cls, name, bases, dict):
     method = dict.get('run_test', BaseUnitTest.run_test)
     for name, args in dict['tests'].items():
         dict['test_%s' % (name, )] = partialmethod(method, dict['func'], args['output'], **args['kwargs'])
     return type(name, bases, dict)