Python 工匠：高效操做文件的三個建議

時間 2019-11-08

原文原文鏈接

前言

這是「Python 工匠」系列的第 11 篇文章。[查看系列全部文章]html

在這個世界上，人們天天都在用 Python 完成着不一樣的工做。而文件操做，則是你們最常須要解決的任務之一。使用 Python，你能夠輕鬆爲他人生成精美的報表，也能夠用短短几行代碼快速解析、整理上萬份數據文件。python

當咱們編寫與文件相關的代碼時，一般會關注這些事情：個人代碼是否是足夠快？個人代碼有沒有事半功倍的完成任務？ 在這篇文章中，我會與你分享與之相關的幾個編程建議。我會向你推薦一個被低估的 Python 標準庫模塊、演示一個讀取大文件的最佳方式、最後再分享我對函數設計的一點思考。git

下面，讓咱們進入第一個「模塊安利」時間吧。github

**注意：**由於不一樣操做系統的文件系統大不相同，本文的主要編寫環境爲 Mac OS/Linux 系統，其中一些代碼可能並不適用於 Windows 系統。編程

建議一：使用 pathlib 模塊

若是你須要在 Python 裏進行文件處理，那麼標準庫中的 os 和 os.path 兄弟倆必定是你沒法避開的兩個模塊。在這兩個模塊裏，有着很是多與文件路徑處理、文件讀寫、文件狀態查看相關的工具函數。bash

讓我用一個例子來展現一下它們的使用場景。有一個目錄裏裝了不少數據文件，可是它們的後綴名並不統一，既有 .txt，又有 .csv。咱們須要把其中以 .txt 結尾的文件都修改成 .csv 後綴名。session

咱們能夠寫出這樣一個函數：編程語言

import os
import os.path


def unify_ext_with_os_path(path):
    """統一目錄下的 .txt 文件名後綴爲 .csv """
    for filename in os.listdir(path):
        basename, ext = os.path.splitext(filename)
        if ext == '.txt':
            abs_filepath = os.path.join(path, filename)
            os.rename(abs_filepath, os.path.join(path, f'{basename}.csv'))
複製代碼

讓咱們看看，上面的代碼一共用到了哪些與文件處理相關的函數：函數

os.listdir(path)：列出 path 目錄下的全部文件*（含文件夾）*
os.path.splitext(filename)：切分文件名裏面的基礎名稱和後綴部分
os.path.join(path, filename)：組合須要操做的文件名爲絕對路徑
os.rename(...)：重命名某個文件

上面的函數雖然能夠完成需求，但說句實話，即便在寫了不少年 Python 代碼後，我依然以爲：這些函數不光很難記，並且最終的成品代碼也不怎麼討人喜歡。工具

使用 pathlib 模塊改寫代碼

爲了讓文件處理變得更簡單，Python 在 3.4 版本引入了一個新的標準庫模塊：pathlib。它基於面向對象思想設計，封裝了很是多與文件操做相關的功能。若是使用它來改寫上面的代碼，結果會大不相同。

使用 pathlib 模塊後的代碼：

from pathlib import Path

def unify_ext_with_pathlib(path):
    for fpath in Path(path).glob('*.txt'):
        fpath.rename(fpath.with_suffix('.csv'))
複製代碼

和舊代碼相比，新函數只須要兩行代碼就完成了工做。而這兩行代碼主要作了這麼幾件事：

首先使用 Path(path) 將字符串路徑轉換爲 Path 對象
調用 .glob('*.txt') 對路徑下全部內容進行模式匹配並以生成器方式返回，結果仍然是 Path 對象，因此咱們能夠接着作後面的操做
使用 .with_suffix('.csv') 直接獲取使用新後綴名的文件全路徑
調用 .rename(target) 完成重命名

相比 os 和 os.path，引入 pathlib 模塊後的代碼明顯更精簡，也更有總體統一感。全部文件相關的操做都是一站式完成。

其餘用法

除此以外，pathlib 模塊還提供了不少有趣的用法。好比使用 / 運算符來組合文件路徑：

# 😑 舊朋友：使用 os.path 模塊
>>> import os.path
>>> os.path.join('/tmp', 'foo.txt')
'/tmp/foo.txt'

# ✨ 新潮流：使用 / 運算符
>>> from pathlib import Path
>>> Path('/tmp') / 'foo.txt'
PosixPath('/tmp/foo.txt')
複製代碼

或者使用 .read_text() 來快速讀取文件內容：

# 標準作法，使用 with open(...) 打開文件
>>> with open('foo.txt') as file:
...     print(file.read())
...
foo

# 使用 pathlib 可讓這件事情變得更簡單
>>> from pathlib import Path
>>> print(Path('foo.txt').read_text())
foo

複製代碼

除了我在文章裏介紹的這些，pathlib 模塊還提供了很是多有用的方法，強烈建議去官方文檔詳細瞭解一下。

若是上面這些都不足以讓你動心，那麼我再多給你一個使用 pathlib 的理由：PEP-519 裏定義了一個專門用於「文件路徑」的新對象協議，這意味着從該 PEP 生效後的 Python 3.6 版本起，pathlib 裏的 Path 對象，能夠和之前絕大多數只接受字符串路徑的標準庫函數兼容使用：

>>> p = Path('/tmp')
# 能夠直接對 Path 類型對象 p 進行 join
>>> os.path.join(p, 'foo.txt')
'/tmp/foo.txt'
複製代碼

因此，無需猶豫，趕忙把 pathlib 模塊用起來吧。

Hint: 若是你使用的是更早的 Python 版本，能夠嘗試安裝 pathlib2 模塊。

建議二：掌握如何流式讀取大文件

幾乎全部人都知道，在 Python 裏讀取文件有一種「標準作法」：首先使用 with open(fine_name) 上下文管理器的方式得到一個文件對象，而後使用 for 循環迭代它，逐行獲取文件裏的內容。

下面是一個使用這種「標準作法」的簡單示例函數：

def count_nine(fname):
    """計算文件裏包含多少個數字 '9' """
    count = 0
    with open(fname) as file:
        for line in file:
            count += line.count('9')
    return count
複製代碼

假如咱們有一個文件 small_file.txt，那麼使用這個函數能夠輕鬆計算出 9 的數量。

# small_file.txt
feiowe9322nasd9233rl
aoeijfiowejf8322kaf9a

# OUTPUT: 3
print(count_nine('small_file.txt'))
複製代碼

爲何這種文件讀取方式會成爲標準？這是由於它有兩個好處：

with 上下文管理器會自動關閉打開的文件描述符
在迭代文件對象時，內容是一行一行返回的，不會佔用太多內存

標準作法的缺點

但這套標準作法並不是沒有缺點。若是被讀取的文件裏，根本就沒有任何換行符，那麼上面的第二個好處就不成立了。當代碼執行到 for line in file 時，line 將會變成一個很是巨大的字符串對象，消耗掉很是可觀的內存。

讓咱們來作個試驗：有一個 5GB 大的文件 big_file.txt，它裏面裝滿了和 small_file.txt 同樣的隨機字符串。只不過它存儲內容的方式稍有不一樣，全部的文本都被放在了同一行裏：

# FILE: big_file.txt
df2if283rkwefh... <剩餘 5GB 大小> ...
複製代碼

若是咱們繼續使用前面的 count_nine 函數去統計這個大文件裏 9 的個數。那麼在個人筆記本上，這個過程會足足花掉 65 秒，並在執行過程當中吃掉機器 2GB 內存 [注1]。

使用 read 方法分塊讀取

爲了解決這個問題，咱們須要暫時把這個「標準作法」放到一邊，使用更底層的 file.read() 方法。與直接循環迭代文件對象不一樣，每次調用 file.read(chunk_size) 會直接返回從當前位置日後讀取 chunk_size 大小的文件內容，沒必要等待任何換行符出現。

因此，若是使用 file.read() 方法，咱們的函數能夠改寫成這樣:

def count_nine_v2(fname):
    """計算文件裏包含多少個數字 '9'，每次讀取 8kb """
    count = 0
    block_size = 1024 * 8
    with open(fname) as fp:
        while True:
            chunk = fp.read(block_size)
            # 當文件沒有更多內容時，read 調用將會返回空字符串 ''
            if not chunk:
                break
            count += chunk.count('9')
    return count
複製代碼

在新函數中，咱們使用了一個 while 循環來讀取文件內容，每次最多讀取 8kb 大小，這樣能夠避免以前須要拼接一個巨大字符串的過程，把內存佔用下降很是多。

利用生成器解耦代碼

假如咱們在討論的不是 Python，而是其餘編程語言。那麼能夠說上面的代碼已經很好了。可是若是你認真分析一下 count_nine_v2 函數，你會發如今循環體內部，存在着兩個獨立的邏輯：數據生成（read 調用與 chunk 判斷） 與 數據消費。而這兩個獨立邏輯被耦合在了一塊兒。

正如我在《編寫地道循環》裏所提到的，爲了提高複用能力，咱們能夠定義一個新的 chunked_file_reader 生成器函數，由它來負責全部與「數據生成」相關的邏輯。這樣 count_nine_v3 裏面的主循環就只須要負責計數便可。

def chunked_file_reader(fp, block_size=1024 * 8):
    """生成器函數：分塊讀取文件內容 """
    while True:
        chunk = fp.read(block_size)
        # 當文件沒有更多內容時，read 調用將會返回空字符串 ''
        if not chunk:
            break
        yield chunk


def count_nine_v3(fname):
    count = 0
    with open(fname) as fp:
        for chunk in chunked_file_reader(fp):
            count += chunk.count('9')
    return count
複製代碼

進行到這一步，代碼彷佛已經沒有優化的空間了，但其實否則。iter(iterable) 是一個用來構造迭代器的內建函數，但它還有一個更少人知道的用法。當咱們使用 iter(callable, sentinel) 的方式調用它時，會返回一個特殊的對象，迭代它將不斷產生可調用對象 callable 的調用結果，直到結果爲 setinel 時，迭代終止。

def chunked_file_reader(file, block_size=1024 * 8):
    """生成器函數：分塊讀取文件內容，使用 iter 函數 """
    # 首先使用 partial(fp.read, block_size) 構造一個新的無需參數的函數
    # 循環將不斷返回 fp.read(block_size) 調用結果，直到其爲 '' 時終止
    for chunk in iter(partial(file.read, block_size), ''):
        yield chunk
複製代碼

最終，只須要兩行代碼，咱們就完成了一個可複用的分塊文件讀取函數。那麼，這個函數在性能方面的表現如何呢？

和一開始的 2GB 內存/耗時 65 秒 相比，使用生成器的版本只須要 7MB 內存 / 12 秒 就能完成計算。效率提高了接近 4 倍，內存佔用更是不到原來的 1%。

建議三：設計接受文件對象的函數

統計完文件裏的「9」以後，讓咱們換一個需求。如今，我想要統計每一個文件裏出現了多少個英文元音字母*（aeiou）*。只要對以前的代碼稍做調整，很快就能夠寫出新函數 count_vowels。

def count_vowels(filename):
    """統計某個文件中，包含元音字母(aeiou)的數量 """
    VOWELS_LETTERS = {'a', 'e', 'i', 'o', 'u'}
    count = 0
    with open(filename, 'r') as fp:
        for line in fp:
            for char in line:
                if char.lower() in VOWELS_LETTERS:
                    count += 1
    return count


# OUTPUT: 16
print(count_vowels('small_file.txt'))
複製代碼

和以前「統計 9」的函數相比，新函數變得稍微複雜了一些。爲了保證程序的正確性，我須要爲它寫一些單元測試。但當我準備寫測試時，卻發現這件事情很是麻煩，主要問題點以下：

函數接收文件路徑做爲參數，因此咱們須要傳遞一個實際存在的文件
爲了準備測試用例，我要麼提供幾個樣板文件，要麼寫一些臨時文件
而文件是否能被正常打開、讀取，也成了咱們須要測試的邊界狀況

**若是，你發現你的函數難以編寫單元測試，那一般意味着你應該改進它的設計。**上面的函數應該如何改進呢？答案是：讓函數依賴「文件對象」而不是文件路徑。

修改後的函數代碼以下：

def count_vowels_v2(fp):
    """統計某個文件中，包含元音字母(aeiou)的數量 """
    VOWELS_LETTERS = {'a', 'e', 'i', 'o', 'u'}
    count = 0
    for line in fp:
        for char in line:
            if char.lower() in VOWELS_LETTERS:
                count += 1
    return count


# 修改函數後，打開文件的職責被移交給了上層函數調用者
with open('small_file.txt') as fp:
    print(count_vowels_v2(fp))
複製代碼

**這個改動帶來的主要變化，在於它提高了函數的適用面。**由於 Python 是「鴨子類型」的，雖然函數須要接受文件對象，但其實咱們能夠把任何實現了文件協議的「類文件對象（file-like object）」傳入 count_vowels_v2 函數中。

而 Python 中有着很是多「類文件對象」。好比 io 模塊內的 StringIO 對象就是其中之一。它是一種基於內存的特殊對象，擁有和文件對象幾乎一致的接口設計。

利用 StringIO，咱們能夠很是方便的爲函數編寫單元測試。

# 注意：如下測試函數須要使用 pytest 執行
import pytest
from io import StringIO


@pytest.mark.parametrize(
    "content,vowels_count", [
        # 使用 pytest 提供的參數化測試工具，定義測試參數列表
        # (文件內容, 期待結果)
        ('', 0),
        ('Hello World!', 3),
        ('HELLO WORLD!', 3),
        ('你好，世界', 0),
    ]
)
def test_count_vowels_v2(content, vowels_count):
    # 利用 StringIO 構造類文件對象 "file"
    file = StringIO(content)
    assert count_vowels_v2(file) == vowels_count
複製代碼

使用 pytest 運行測試能夠發現，函數能夠經過全部的用例：

❯ pytest vowels_counter.py
====== test session starts ======
collected 4 items

vowels_counter.py ... [100%]

====== 4 passed in 0.06 seconds ======
複製代碼

而讓編寫單元測試變得更簡單，並不是修改函數依賴後的惟一好處。除了 StringIO 外，subprocess 模塊調用系統命令時用來存儲標準輸出的 PIPE 對象，也是一種「類文件對象」。這意味着咱們能夠直接把某個命令的輸出傳遞給 count_vowels_v2 函數來計算元音字母數：

import subprocess

# 統計 /tmp 下面全部一級子文件名（目錄名）有多少元音字母
p = subprocess.Popen(['ls', '/tmp'], stdout=subprocess.PIPE, encoding='utf-8')

# p.stdout 是一個流式類文件對象，能夠直接傳入函數
# OUTPUT: 42
print(count_vowels_v2(p.stdout))
複製代碼

正如以前所說，將函數參數修改成「文件對象」，最大的好處是提升了函數的 適用面 和 可組合性。經過依賴更爲抽象的「類文件對象」而非文件路徑，給函數的使用方式開啓了更多可能，StringIO、PIPE 以及任何其餘知足協議的對象均可以成爲函數的客戶。

不過，這樣的改造並不是毫完好點，它也會給調用方帶來一些不便。假如調用方就是想要使用文件路徑，那麼就必須得自行處理文件的打開操做。

如何編寫兼容兩者的函數

有沒有辦法即擁有「接受文件對象」的靈活性，又能讓傳遞文件路徑的調用方更方便？答案是：有，並且標準庫中就有這樣的例子。

打開標準庫裏的 xml.etree.ElementTree 模塊，翻開裏面的 ElementTree.parse 方法。你會發現這個方法便可以使用文件對象調用，也接受字符串的文件路徑。而它實現這一點的手法也很是簡單易懂：

def parse(self, source, parser=None):
    """*source* is a file name or file object, *parser* is an optional parser """
    close_source = False
    # 經過判斷 source 是否有 "read" 屬性來斷定它是否是「類文件對象」
    # 若是不是，那麼調用 open 函數打開它並負擔起在函數末尾關閉它的責任
    if not hasattr(source, "read"):
        source = open(source, "rb")
        close_source = True
複製代碼