最近在试图实现 pdir2 彻底 disable color 的 feature, 让它和其他 cli 的做法一样:在 stdout 是 TTY 的时候默认开启带有颜色的输出,在 stdout 不是 TTY 的时候不输出颜色,即没有颜色相关的 escape code. 这个 pull request 在这里。
实现比较简单,在非 TTY 的情况下,新建了一个 Fake 的 Color Render,没有做任何渲染,直接输出。
但是在测试的时候遇到了难题:之前的测试都是按照默认输出颜色来写的,而 pytest 运行的时候,显然是没有 TTY 的,我的代码改变了这个行为,导致 pytest 运行测试的时候,颜色都消失了。
我试图用 patch 来解决这个问题,设定一个全局的 fixture,让 pytest 运行的时候,sys.stdout.isatty()
返回的是 True, 这样所有之前的 test case 都可以依然 pass. 然后就发现了一个非常难解决的问题。
使用 sys.stdout.isatty()
的地方,是在一个 module 的全局的地方,判断之后设置了一个全局变量,类似如下:
1 2 |
if sys.stdout.isatty(): use_color = True |
而这个 module 在 pytest collecting tests 的时候就已经 import 了,所以这个 use_color
缓存在了 module 的全局中,即使我去 patch isatty()
, 实际上也不会调用到了。
于是我想到删除这个缓存。在 Python3 里面做这件事很简单,用 importlib.reload("pdir")
就可以了。
然而又出现了一个棘手的问题:pdir
这个 module,其实在 import 之后就已经不存在了。作者为了想让用户这么使用:import pdir; pdir(foo)
而不需要 from pdir import pdir; pdir(foo)
,用了一个 trick:即,在 pdir.__init__
中,直接将 sys.module['pdir']
替换掉了: sys.modules[__name__] = PrettyDir
,这样的好处是:import 进来的不再是一个 module,而是一个 class,直接可以调用了。但是坏处是,我们再也找不到 pdir 这个 module 了,也就无法使用 importlib 进行 reload.
为了能够让它在 patch 之后进行 reload,我尝试了很多 hack,比如直接 patch 它的全局变量,发现会失败,因为 patch
也找不到 target 了,因为 patch
也找不到 pdir 这个 module,会提示 class 没有你要 patch 的这个属性;另外尝试过,在源代码中,sys.modules[__name__] = PrettyDir
替换之前先保存原来的 module 到一个新的名字,发现也不行,这个 import 机制貌似必须让 module 的名字和 module 对应。
睡了一觉之后,我又在思考为什么 pytest 在收集测试的时候就运行了 module 的 init 呢?为什么不让我先 patch 好,它再去 import?
又仔细看了代码,发现,有一些 test.py
文件,在文件的开头就 import pdir
了。这时候,无论我怎么 patch,后面运行测试的时候都会使用已经 init 好的。
所以要解决这个问题,其实很简单:
- test 文件不能再任何地方
import pdir
,必须要在 test case 里面进行 import,这样,收集测试的时候不会初始化 module - 然后我设置一个全局的 fixture 去 patch
isatty()
。这样,执行的逻辑就变成了:收集测试(没有 pdir init)-> 全局 fixture 执行,isatty()
patch 为 True -> 执行测试 -> 测试内部 import pdir -> pdir 认为 stdout 是一个 TTY
这样还有一个不好的地方,就是全局初始化好了,测试中就无法再测试不是 TTY 的逻辑了。
最后在代码中看到一段 sys.modules
的删除逻辑,好像作者也遇到过类似全局变量在测试中需要重新初始化的需求。把这段代码放到 fixture,发现居然神奇的工作了。原理很简单,就是我 patch 了之后,需要删除所有 pdir 的缓存,这样 import 的时候,就会重新 init 一遍。需要注意的是,不能只删除 pdir
, pdir.*
都需要删除。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
def remove_module_cache(): for module in ( 'pdir', 'pdir.api', 'pdir.constants', 'pdir.format', 'pdir.configuration', 'pdir.color', 'pdir.utils', ): try: del modules[module] except KeyError: pass @pytest.fixture(scope="session") def tty(): with patch("sys.stdout.isatty") as faketty: faketty.return_value = True remove_module_cache() yield |
这样,只需要在测试 TTY 的时候,使用 tty
这个 fixture,不使用的话就默认不是一个 TTY。
1 2 3 4 5 6 7 |
def test_is_tty(tty): import pdir ... def test_not_tty(): import pdir ... |
要解决这个问题, 还有其他一些可能的思路有:
- 让 pytest 为每一个测试(或者测试文件)重新开启一个 python 解释器,这样就完全干净了。但是没看到 pytest 有这样的 feature
- 减少全局变量的使用,每一次调用都判断是否是 TTY 的逻辑,这样就是为了测试去修改原来的逻辑了,不太喜欢