【Python】キーでグルーピングする|itertools.groupby

【Python】キーでグルーピングする

やること

キーでグルーピングします。
例として、都道府県と市区町村が1:1のリストを、都道府県ごとにまとめます。

Before

[
    {'prefecture' : '北海道', 'city' : '札幌市'},
    {'prefecture' : '北海道', 'city' : '函館市'},
    {'prefecture' : '北海道', 'city' : '小樽市'},
    {'prefecture' : '北海道', 'city' : '旭川市'},

    {'prefecture' : '青森県', 'city' : '青森市'},
    {'prefecture' : '青森県', 'city' : '弘前市'},
    {'prefecture' : '青森県', 'city' : '八戸市'}
]

After

[
    Prefecture(prefecture='北海道', cites=['札幌市', '函館市', '小樽市', '旭川市']),
    Prefecture(prefecture='青森県', cites=['青森市', '弘前市', '八戸市'])
]

コード

import dataclasses
import itertools
from operator import itemgetter


@dataclasses.dataclass
class Prefecture:
    """都道府県データ
    """
    # 都道府県
    prefecture: str
    # 市区町村リスト
    cites: list


class PrefectureList:
    """都道府県リスト
    """
    # 都道府県リスト
    Prefectures: list

    def groupby_prefecture(self, areas):
        """都道府県ごとリスト生成

        Args:
            areas (list): 都道府県リスト。
                          [{'prefecture':'都道府県名', 'city':'市区町村名'}]
        Returns:
            result: 都道府県ごとにグルーピングしたリスト
        """
        # prefectureをキーにソートする
        key = 'prefecture'
        areas.sort(key=itemgetter(key))

        result = []
        for key, group in itertools.groupby(areas, key=itemgetter(key)):
            # 市区町村をリストにする
            cities = [member['city'] for member in group]
            # Prefectureのインスタンスを生成し、データを格納する。
            pref = Prefecture(key, cities)
            # 返却用のリストに、生成したPrefectureのインスタンスを追加する。
            result.append(pref)

        # 生成したリストを返却する。
        return result

解説

PrefectureクラスはC言語で言う構造体として定義しています。
自分が元々C言語やJavaなので、データクラスを使うほうが馴染みやすいという理由です。

itertools.groupby

グループ化はitertools.groupbyを使います。

itertools.groupbyは事前にソートされていることが前提なので、33行目でソートしています。

グループ化したら、ループで取り出して、Prefectureに格納して、リストにして行きます。

参考 itertools

テストコード

テストコードもつけておきますね。
前述のとおり自分が元々Javaの人なので、JUnitの思想を受け継ぐunittestが使いやすくて好きです。

test_prefecture.py

shortDescription()をオーバーライドして、DocStringをすべて出力させているのは、エビデンスの整理をやりやすくするためです。
邪魔であれば削除してもかまいません。

import unittest

import tests.testutil as testutil

from prefecture import Prefecture, PrefectureList


class Test〇〇(unittest.TestCase):
    def setUp(self) -> None:
        """
        テストメソッド実行前の処理
        """
        self.target = PrefectureList()

    def tearDown(self) -> None:
        """
        テストメソッド実行後の処理
        """
        return super().tearDown()

    def shortDescription(self):
        """Superクラスのメソッドを上書き
           DocStringを改行で区切らず、すべて出力させる。

        Returns:
            doc: DocString
        """
        doc = self._testMethodDoc
        return doc if doc else None

    def test_prefecture_01(self):
        """
        ----------------------------------------------------------------
        ◆ prefectureをキーにグルーピングされること
          1件
        ----------------------------------------------------------------
        """
        # テストデータ
        prefectures = [
            {'prefecture' : '北海道', 'city' : '札幌市'},
        ]

        expect = [
            Prefecture('北海道', ['札幌市'])
        ]

        # テスト対象実行
        result = self.target.groupby_prefecture(prefectures)
        # テスト内容
        testutil.print_result(prefectures, result)
        # 検証
        self.assertEqual(result, expect)

    def test_prefecture_02(self):
        """
        ----------------------------------------------------------------
        ◆ prefectureをキーにグルーピングされること
          複数件
        ----------------------------------------------------------------
        """
        # DocStringを出力
        # self.print_docstring(sys._getframe().f_code.co_name)

        # テストデータ
        prefectures = [
            {'prefecture' : '北海道', 'city' : '札幌市'},
            {'prefecture' : '北海道', 'city' : '函館市'},
            {'prefecture' : '北海道', 'city' : '小樽市'},
            {'prefecture' : '北海道', 'city' : '旭川市'},

            {'prefecture' : '青森県', 'city' : '青森市'},
            {'prefecture' : '青森県', 'city' : '弘前市'},
            {'prefecture' : '青森県', 'city' : '八戸市'}
        ]

        expect = [
            Prefecture('北海道', ['札幌市', '函館市', '小樽市', '旭川市']),
            Prefecture('青森県', ['青森市', '弘前市', '八戸市'])
        ]

        # テスト対象実行
        result = self.target.groupby_prefecture(prefectures)
        # テスト内容
        testutil.print_result(prefectures, result)
        # 検証
        self.assertEqual(result, expect)

testutil.py

テストデータの生成や、エビデンス整理しやすいように使うメソッドを書くモジュールを別モジュールにしてます。
わたしの好みで分けていますが、たぶんあまり一般的なやり方ではないと思うので、真似しないほうが良いんじゃないかな。

    def print_result(before, after):
    """Before, Afterを出力
    Args:
        before (object): Before
        after (object): After
    """
    print('Before'.ljust(80, '-'), '\n', before, '\n')
    print('After'.ljust(80, '-'), '\n', after, '\n')

詳細設計書

詳細設計書やプログラム設計書不要論もありますが、書けと言うところは今でもそれなりにあるので、設計書も置いておきますね。

    # 都道府県ごとリスト生成(groupby_prefecture)

    ## 概要
    
    都道府県:市区町村のリストを、都道府県ごとにグルーピングする。
    
    例)
    ```python
    # Before
    [
        {'prefecture' : '北海道', 'city' : '札幌市'},
        {'prefecture' : '北海道', 'city' : '函館市'},
        {'prefecture' : '北海道', 'city' : '小樽市'},
        {'prefecture' : '北海道', 'city' : '旭川市'},
    
        {'prefecture' : '青森県', 'city' : '青森市'},
        {'prefecture' : '青森県', 'city' : '弘前市'},
        {'prefecture' : '青森県', 'city' : '八戸市'}
    ]
    ```
    
    ```python
    # After
    [
        Prefecture(prefecture='北海道', cites=['札幌市', '函館市', '小樽市', '旭川市']),
        Prefecture(prefecture='青森県', cites=['青森市', '弘前市', '八戸市'])
    ]
    ```
    
    ## 引数
    
    | 引数 | I/O | 内容 | 備考 |
    | -- | -- | -- | -- |
    | areas | In | 都道府県リスト | [{'都道府県':'市区町村'},{'都道府県':'市区町村'}…] |
    
    ## 戻り値
    
    | 戻り値 | 備考 |
    | -- | -- |
    | 都道府県リスト | [Prefecture]() |
    
    
    ## ワークフロー
    
    1. 引数.都道府県リスト(areas)を、「prefecture」をキーにソートする。
    2. ソートしたリストを、「prefecture」ごとにループする。
       1. 市区町村をリストにする。
       2. Prefectureのインスタンスを生成し、データを格納する。
       3. 返却用のリストに、生成したPrefectureのインスタンスを追加する。
    3. 生成したリストを返却する。
    

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です

CAPTCHA