一日一技：使用装饰器简化大量if判断（二）

发表于 更新于 分类于 阅读次数： 本文字数： 1.6k 阅读时长 ≈ 5 分钟

在之前的文章：一日一技：使用装饰器简化大量 if…elif…代码发布以后，有很多同学说想看后续，如何在装饰器中表示大于小于。甚至有同学每周来催一次稿：

于是，今天我们就来看看大于小于应该怎么来判断。为了实现我们今天的目标，有两个前置知识需要掌握，一个是Python自带的operator模块，另一个是偏函数。

2 > 1还有另一种写法？

当我们要表达大于这个意思的时候，你想到的肯定是大于符号>。所以2大于1，肯定写作2 > 1。这看起来是很正常的事情。现在，如果我让你不准使用大于符号>，怎么表示大于？

实际上，在Python里面，除了>外，还有一种写法，就是使用自带的operator模块：

1
2
3
import operator

operator.gt(2, 1)

其中的.gt(参数1, 参数2)就表示参数1 > 参数2。如果成立，返回True，否则返回False。

类似的还有：

  • 大于等于：operator.ge
  • 小于：operator.lt
  • 小于等于：operator.le
  • 不等于：operator.ne
  • 等于：operator.eq

因此，下面两个写法是等价的：

1
2
if a <= b:
    print('成功')
1
2
if operator.le(a, b):
    print('成功')

偏函数

我在很久以前的公众号文章里面已经介绍过偏函数了：偏函数：在Python中设定默认参数的另一种办法。因此本文就不再讲它的基础用法了，大家点击链接去看那篇文章就可以掌握。

为什么我们需要偏函数呢？这是因为我们今天要做的事情，它需要给函数先传一半的参数，另一半的参数要在未来才能传入。例如，循环等待用户输入数字，如果其中一次输入的数字大于等于5，就打印你好世界。

如果不知道偏函数，你可能是这样写的：

1
2
3
4
while True:
    num = int(input('请输入数字：'))
    if num >= 5:
        print('你好世界')

有了偏函数以后，你的写法是这样的：

1
2
3
4
5
6
7
import operator
from functools import partial
ge_5 = partial(operator.le, 5)
while True:
    num = int(input('请输入数字：'))
    if ge_5(num):
        print('你好世界')

特别注意，这里我在偏函数中传入的第一个参数是operator.le：小于。因为operator.xx表示第一个参数对第二个参数的比较，所以x >= 5 就相当于5 <= x 也就是operator.le(5, x)。

在装饰器中实现大小比较

前置知识掌握以后，我们就能看如何在装饰器里面实现大小比较。在第一篇文章中，我们只实现了参数等于，它的原理是：

1
2
3
4
5
6
7
8
9
10
def register(value):
        def wrap(func):
            if value in registry:
                raise ValueError(
                    f'@value_dispatch: there is already a handler '
                    f'registered for {value!r}'
                )
            registry[value] = func
            return func
        return wrap

register只接收了一个位置参数value。但实际上，我们还可以通过修改这段注册的代码，实现如下的效果：

1
2
3
4
5
6
7
@get_discount.register(3, op='gt')
def parse_level_gt3(level):
    print('等级大于3')

@get_discount.register(3, op='le')
def parse_level_le3(level):
    print('等级小于等于3')

有同学问，有没有可能实现这样的写法呢：

1
2
3
@get_discount.register(2, le=3)
def parse_level_gt3(level):
    print('等级为2')

我觉得这样写是没有什么必要的。因为register()里面，多个参数之间的关系是且。那么只有两种情况，要么，就等于这个数，例如@get_discount.register(2, le=3)，既要等于2，又要小于等于3，那显然就等于2。不需要写这个le=3。要么，就不存在结果，例如@get_discount.register(2, gt=3)，既要等于2，又要大于3，显然下面被装饰的函数永远不会执行。因为找不到这个数。

因此，我们的装饰器函数就可以做如下修改：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import functools
import operator

def value_dispatch(func):

    registry_eq = {}
    registry_other = {}
    key_op_map = {}

    @functools.wraps(func)
    def wrapper(arg0, *args, **kwargs):
        if arg0 in registry_eq:
            delegate = registry_eq[arg0]
            return delegate(arg0, *args, **kwargs)
        else:
            for key, op in key_op_map.items():
                if op(arg0):
                    delegate = registry_other[key]
                    return delegate(arg0, *args, **kwargs)
        return func(arg0, *args, **kwargs)

    def register(value, op='eq'):
        if op == 'eq':
            def wrap(func):
                if value in registry_eq:
                    raise ValueError(
                        f'@value_dispatch: there is already a handler '
                        f'registered for {value!r}'
                    )
                registry_eq[value] = func
                return func
            return wrap
        else:
            if op == 'gt':
                op_func = functools.partial(operator.lt, value)
            elif op == 'ge':
                op_func = functools.partial(operator.le, value)
            elif op == 'lt':
                op_func = functools.partial(operator.gt, value)
            elif op == 'le':
                op_func = functools.partial(operator.ge, value)
            else:
                raise ValueError('op 参数只能是：gt/ge/lt/le之一')
            key = f'{op}_{value}'
            key_op_map[key] = op_func
            def wrap(func):
                if key in registry_other:
                    raise ValueError(
                        f'@value_dispatch: there is already a handler '
                        f'registered for {key!r}'
                    )
                registry_other[key] = func
                return func
            return wrap
            
    
    wrapper.register = register
    return wrapper

它的使用方法还是跟以前一样，先定义默认的函数逻辑：

1
2
3
@value_dispatch
def get_discount(level):
    return '等级错误'

如果定义相等的逻辑，写法跟以前完全一样：

1
2
3
4
5
@get_discount.register(1)
def parse_level_1(level):
    "大量计算代码"
    discount = 0.1
    return discount

如果要定义不等于逻辑，就在.register()中添加一个参数op：

1
2
3
4
@get_discount.register(2, op='gt')
def parse_level_gt2(level):
    discount = 1
    return 1

运行效果如下图所示：

由于我们定义了大于2时，始终返回1，所以可以看到get_discount(6)和get_discount(10)返回的都是1.

由于我们只定义了等于1和大于2的逻辑，所以当传入的参数为2时，就返回等级错误.

到这里，本文要讲的内容就结束了。但最后还是要考大家3个问题：

如果不使用偏函数和operator模块，你会怎么做

你可以试一试在不实用偏函数和operator的情况下，实现这个需求。

如果定义的条件有重叠怎么办？

例如对于下面的两个函数：

1
2
3
4
5
6
7
8
9
@get_discount.register(2, op='gt')
def parse_level_gt2(level):
    discount = 1
    return discount

@get_discount.register(10, op='gt')
def parse_level_gt2(level):
    discount = 100
    return discount

当level的值是20的时候，同时满足两个条件，应该运行哪一个呢？

如何定义区间？

怎么实现这样的功能：

1
2
3
4
5
@get_discount.register(ge=2, lt=5)
def parse_level_between2_5(level):
    print('等级2<=level<5')
    discount = 0.5
    return discount

如果区间存在全包含、部分包含应该运行哪个函数？例如：

1
2
3
4
5
6
7
8
@get_discount.register(ge=2, lt=00)
...

@get_discount.register(ge=20, lt=50)
...

@get_discount.register(ge=80, lt=200)
...

请大家把你对这两个问题的答案回答在评论区里面。提示（想清楚什么是真需求，什么是伪需求，再考虑怎么解决）