Python 上下文管理及 with 语句的实用技巧-工具盒子

Python 有很多魔法方法，本文记录一下可以自定义 with 语句的上下文管理器所使用到的两个魔法方法，也就是 __enter__ 和 __exit__ 方法的实用性。

自定义上下文管理类 {#自定义上下文管理类}

最常见的 with 语句就是 open 函数了，这里不做解释，直接来看一个自定义类的例子。

class TestHandler():
    def __init__(self):
        pass

    <span class="k">def</span> <span class="fm">__enter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span>

    <span class="k">def</span> <span class="fm">__exit__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">exc_type</span><span class="p">,</span> <span class="n">exc_val</span><span class="p">,</span> <span class="n">exc_tb</span><span class="p">):</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">'exc_type:'</span><span class="p">,</span> <span class="n">exc_type</span><span class="p">)</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">'exc_val:'</span><span class="p">,</span> <span class="n">exc_val</span><span class="p">)</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">'exc_tb:'</span><span class="p">,</span> <span class="n">exc_tb</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">func</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="nb">print</span><span class="p">(</span><span class="mi">1</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">bad_func</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">'a'</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>

上面定义了一个类，这个类定义了两个打印值的方法，其中一个方法会报错，同时该类使用到了两个魔法方法，有了这两个方法，这个类就可以使用 with 语句来进行调用，来看看调用正常函数的结果：

with TestHandler() as t:
    t.func()

结果如下：

2
exc_type: None
exc_val: None
exc_tb: None

再来看看调用报错函数的结果

with TestHandler() as t:
    t.bad_func()

exc_type: <class 'TypeError'>
exc_val: Can't convert 'int' object to str implicitly
exc_tb: <traceback object at 0x0000021CEB484B08>
Traceback (most recent call last):
  File "D:/Mycode/TestCase/mark.py", line 23, in <module>
    t.bad_func()
  File "D:/Mycode/TestCase/mark.py", line 17, in bad_func
    print('a' + 1)
TypeError: Can't convert 'int' object to str implicitly

从上面两次调用，可以看到，__exit__ 函数里面的三个参数（定义函数的时候默认会要求加入）分别代表了报错类型、报错原因、报错追溯，只有当 with 语句调用报错时候，这三个参数才有值，否则就是 None，看到这里，你是否能够想到什么？可以利用这三个参数进行异常判断和处理。

上下文管理实用性 {#上下文管理实用性}

已经知道如何定义 with 语句了，也知道遇到异常会出现什么，那么现在来看看自定义 with 语句的使用场景有哪些。

数据库连接操作 {#数据库连接操作}

with 语句比较适合的场景是打开->操作->关闭，在我们常用的除了文件操作外，还有数据库操作、SSH 操作会涉及这个过程，所以，直接看看这两个操作的例子。

import sqlite3

class DBHandler():


    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">database</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">database</span> <span class="o">=</span> <span class="n">database</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">conn</span> <span class="o">=</span> <span class="n">sqlite3</span><span class="o">.</span><span class="n">connect</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">database</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cursor</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">conn</span><span class="o">.</span><span class="n">cursor</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__enter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">cursor</span>

    <span class="k">def</span> <span class="fm">__exit__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">exc_type</span><span class="p">,</span> <span class="n">exc_val</span><span class="p">,</span> <span class="n">exc_tb</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">exc_type</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">conn</span><span class="o">.</span><span class="n">commit</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">conn</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>

上面这个关于数据操作的例子就很典型，它包括了数据库连接、数据库操作（with 语句之后）、异常处理、数据库关闭连接等操作。

看一下 with 语句的使用：

with DBHandler(database) as db:
    db.executescript(create_sql)

是不是非常的方便，当然，如果再结合 try 语句来进行连接操作，就更安全可靠。

SSH 连接操作 {#ssh-连接操作}

再来看看 SSH 的操作例子

import paramiko

class SSHClient:
def init(self, hostname, username, password, port=22):
self.hostname = hostname
self.username = username
self.password = password
self.port = port
self.client = None


    <span class="k">def</span> <span class="nf">connect</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span> <span class="o">=</span> <span class="n">paramiko</span><span class="o">.</span><span class="n">SSHClient</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">set_missing_host_key_policy</span><span class="p">(</span><span class="n">paramiko</span><span class="o">.</span><span class="n">AutoAddPolicy</span><span class="p">())</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">connect</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">hostname</span><span class="p">,</span> <span class="n">port</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">port</span><span class="p">,</span> <span class="n">username</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">username</span><span class="p">,</span> <span class="n">password</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">password</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">execute_command</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">command</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">"SSH client is not connected"</span><span class="p">)</span>
        <span class="n">stdin</span><span class="p">,</span> <span class="n">stdout</span><span class="p">,</span> <span class="n">stderr</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">exec_command</span><span class="p">(</span><span class="n">command</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">stdout</span><span class="o">.</span><span class="n">read</span><span class="p">()</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="s1">'utf-8'</span><span class="p">),</span> <span class="n">stderr</span><span class="o">.</span><span class="n">read</span><span class="p">()</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="s1">'utf-8'</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">close</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="fm">__enter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">connect</span><span class="p">()</span>
        <span class="k">return</span> <span class="bp">self</span>

    <span class="k">def</span> <span class="fm">__exit__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">exc_type</span><span class="p">,</span> <span class="n">exc_value</span><span class="p">,</span> <span class="n">traceback</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">exc_type</span><span class="p">:</span>
            <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">"An exception occurred: </span><span class="si">{</span><span class="n">exc_value</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
        <span class="k">return</span> <span class="kc">False</span>  <span class="c1"># 不处理异常，异常会被重新抛出</span>




# 使用示例


# 使用 with 语句
with SSHClient('hostname', 'username', 'password') as ssh:
stdout, stderr = ssh.execute_command('ls -l')
print("STDOUT:", stdout)
print("STDERR:", stderr)


# 交互式执行命令
with SSHClient('hostname', 'username', 'password') as ssh:
shell = ssh.client.invoke_shell()
shell.send('ifconfig\n')
time.sleep(1)
output = shell.recv(1024).decode()
print(output)

# 直接连接`
`ssh` `=` `SSHClient('hostname',` `'username',` `'password')`
`ssh.connect()`
`stdout,` `stderr` `=` `ssh.execute_command('ls -l')`
`print("STDOUT:",` `stdout)`
`print("STDERR:",` `stderr)`
`ssh.close()`
`

很明显，上面的自定义类 with 语句返回的是一个 SSHClient 对象，所以使用时直接按照这个对象的方法调用即可，调用结束会自动断开连接。

Telnet 连接操作 {#telnet-连接操作}

下面是封装的一个Telnet客户端类，以实现与Telnet协议的连接、执行命令和安全关闭连接的功能

import telnetlib
import time

class TelnetClient:
def init(self, hostname, username, password, port=23):
self.hostname = hostname
self.username = username
self.password = password
self.port = port
self.client = None


    <span class="k">def</span> <span class="nf">connect</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span> <span class="o">=</span> <span class="n">telnetlib</span><span class="o">.</span><span class="n">Telnet</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">hostname</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">port</span><span class="p">,</span> <span class="n">timeout</span><span class="o">=</span><span class="mi">10</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_login</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_login</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">read_until</span><span class="p">(</span><span class="sa">b</span><span class="s1">'Username: '</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">username</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s1">'ascii'</span><span class="p">)</span> <span class="o">+</span> <span class="sa">b</span><span class="s1">'</span><span class="se">\n</span><span class="s1">'</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">read_until</span><span class="p">(</span><span class="sa">b</span><span class="s1">'Password: '</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">password</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s1">'ascii'</span><span class="p">)</span> <span class="o">+</span> <span class="sa">b</span><span class="s1">'</span><span class="se">\n</span><span class="s1">'</span><span class="p">)</span>
        <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>  <span class="c1"># 等待登录完成</span>

    <span class="k">def</span> <span class="nf">execute_command</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">command</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">"Telnet client is not connected"</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">command</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s1">'ascii'</span><span class="p">)</span> <span class="o">+</span> <span class="sa">b</span><span class="s1">'</span><span class="se">\n</span><span class="s1">'</span><span class="p">)</span>
        <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>  <span class="c1"># 等待命令执行</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">read_very_eager</span><span class="p">()</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="s1">'ascii'</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">close</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="sa">b</span><span class="s1">'exit</span><span class="se">\n</span><span class="s1">'</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">client</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="fm">__enter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">connect</span><span class="p">()</span>
        <span class="k">return</span> <span class="bp">self</span>

    <span class="k">def</span> <span class="fm">__exit__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">exc_type</span><span class="p">,</span> <span class="n">exc_value</span><span class="p">,</span> <span class="n">traceback</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">exc_type</span><span class="p">:</span>
            <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">"An exception occurred: </span><span class="si">{</span><span class="n">exc_value</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span>
        <span class="k">return</span> <span class="kc">False</span>  <span class="c1"># 不处理异常，异常会被重新抛出</span>




# 使用示例
if name == "main":
with TelnetClient('192.168.1.1', 'admin', 'password') as telnet:
output = telnet.execute_command('show ip interface brief')
print("Command Output:\n", output)


    <span class="c1"># 交互式执行命令</span>
    <span class="k">with</span> <span class="n">TelnetClient</span><span class="p">(</span><span class="s1">'192.168.1.1'</span><span class="p">,</span> <span class="s1">'admin'</span><span class="p">,</span> <span class="s1">'password'</span><span class="p">)</span> <span class="k">as</span> <span class="n">telnet</span><span class="p">:</span>
        <span class="n">shell</span> <span class="o">=</span> <span class="n">telnet</span><span class="o">.</span><span class="n">client</span>  <span class="c1"># 获取Telnet客户端</span>
        <span class="n">shell</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="sa">b</span><span class="s1">'ifconfig</span><span class="se">\n</span><span class="s1">'</span><span class="p">)</span>  <span class="c1"># 执行命令</span>
        <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">shell</span><span class="o">.</span><span class="n">read_very_eager</span><span class="p">()</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="s1">'ascii'</span><span class="p">)</span>  <span class="c1"># 读取输出</span>
        <span class="nb">print</span><span class="p">(</span><span class="s2">"Interactive Command Output:</span><span class="se">\n</span><span class="s2">"</span><span class="p">,</span> <span class="n">output</span><span class="p">)</span>

    <span class="c1"># 直接连接</span>
    <span class="n">telnet</span> <span class="o">=</span> <span class="n">TelnetClient</span><span class="p">(</span><span class="s1">'192.168.1.1'</span><span class="p">,</span> <span class="s1">'admin'</span><span class="p">,</span> <span class="s1">'password'</span><span class="p">)</span>
    <span class="n">telnet</span><span class="o">.</span><span class="n">connect</span><span class="p">()</span>
    <span class="n">output</span> <span class="o">=</span> <span class="n">telnet</span><span class="o">.</span><span class="n">execute_command</span><span class="p">(</span><span class="s1">'show version'</span><span class="p">)</span>
    <span class="nb">print</span><span class="p">(</span><span class="s2">"Command Output:</span><span class="se">\n</span><span class="s2">"</span><span class="p">,</span> <span class="n">output</span><span class="p">)</span>
    <span class="n">telnet</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>

总结：自定义 with 语句简单理解就是非常适合一些"有始有终"的场景，通过自定义上下文管理器，可以把一些需要重复执行的固定操作简化，只需要关注特定的操作本身。