Enum&dict継承で型安全にしてみた

Author: fono09

library/omikuji.py

@@ -4,7 +4,7 @@
 おみくじを返す
 """
 
-from typing import Hashable, Tuple
+from typing import Tuple, TypeVar
 from dataclasses import dataclass
 from random import choices
 
@@ -26,66 +26,18 @@ def __post_init__(self):
         assert self.message != ''
 
 
-@dataclass
-class Omikuji:
-    """
-    おみくじのコアロジック
-    ガチャではないので排出率を公開するメソッドはあえて実装されていない
-    """
-    entries: dict[Hashable, OmikujiResult]
-
-    def draw(self) -> Tuple[Hashable, OmikujiResult]:
-        """
-        おみくじを引く
-        """
-        return choices(
-            population=list(self.entries.items()),
-            weights=list(
-                map(lambda entry: entry.appearance, self.entries.values())),
-            k=1
-        )[0]
-
+TOmikujiEnum = TypeVar('TOmikujiEnum')
+OmikujiResults = dict[TOmikujiEnum, OmikujiResult]
 
-# 以下おみくじの設定
 
-omikuji = Omikuji({
-    'DAI_KICHI': OmikujiResult(
-        200,
-        ":tada: 大吉 何でもうまくいく!!気がする!!"
-    ),
-    'CHU_KICHI': OmikujiResult(
-        2000,
-        ":smile: 中吉 そこそこうまくいくかも!?"
-    ),
-    'SHO_KICHI': OmikujiResult(
-        3800,
-        ":smily: 小吉 なんとなくうまくいくかも!?"
-    ),
-
-    'KICHI': OmikujiResult(
-        3000,
-        ":smirk: 吉 まあうまくいくかも!?"
-    ),
-    'HATO_KICHI': OmikujiResult(
-        900,
-        ":dove_of_peace: 鳩吉 お前が鳩になる番だ!!羽ばたけ!!!飛べ!!!!唐揚げになれ!!!!!"
-    ),
-
-    'KYO': OmikujiResult(
-        75,
-        ":cry: 凶 ちょっと慎重にいったほうがいいかも……"
-    ),
-
-    'DAI_KYO': OmikujiResult(
-        25,
-        ":crying_cat_face: 大凶 そういう時もあります……猫になって耐えましょう"
-    ),
-
-})
-
-
-def draw() -> str:
+def draw(entries: OmikujiResults) -> Tuple[TOmikujiEnum, OmikujiResult]:
     """
-    おみくじ抽選
+    おみくじを引く
     """
-    return omikuji.draw()[1].message
+
+    return choices(
+        population=list(entries.items()),
+        weights=list(
+            map(lambda entry: entry.appearance, entries.values())),
+        k=1
+    )[0]

plugins/hato.py

@@ -6,6 +6,7 @@
 import json
 import os
 import re
+from enum import Enum, auto
 from logging import getLogger
 from tempfile import NamedTemporaryFile
 from typing import List
@@ -22,7 +23,7 @@
 from library.hatokaraage import hato_ha_karaage
 from library.clientclass import BaseClient
 from library.jma_amesh import jma_amesh
-from library.omikuji import draw as omikuji_draw
+from library.omikuji import OmikujiResult, OmikujiResults, draw as omikuji_draw
 logger = getLogger(__name__)
 
 
@@ -277,8 +278,63 @@ def yoshiyoshi(client: BaseClient):
     client.post('よしよし')
 
 
+# 以下おみくじの設定
+# Refer: dev-hato/hato-bot#876
+class OmikujiEnum(Enum):
+    """
+    おみくじの結果一覧
+    """
+    DAI_KICHI = auto()
+    CHU_KICHI = auto()
+    SHO_KICHI = auto()
+    KICHI = auto()
+    SUE_KICHI = auto()
+    AGE_KICHI = auto()
+    KYO = auto()
+    DAI_KYO = auto()
+
+
+omikuji_results = OmikujiResults({
+    OmikujiEnum.DAI_KICHI: OmikujiResult(
+        12,
+        ":tada: 大吉 何でもうまくいく!!気がする!!"
+    ),
+    OmikujiEnum.KICHI: OmikujiResult(
+        100,
+        ":smirk: 吉 まあうまくいくかも!?"
+    ),
+    OmikujiEnum.CHU_KICHI: OmikujiResult(
+        100,
+        ":smile: 中吉 そこそこうまくいくかも!?"
+    ),
+    OmikujiEnum.SHO_KICHI: OmikujiResult(
+        100,
+        ":smily: 小吉 なんとなくうまくいくかも!?"
+    ),
+    OmikujiEnum.SUE_KICHI: OmikujiResult(
+        37,
+        ":expressionless: 末吉 まあ多分うまくいくかもね……!?"
+    ),
+    OmikujiEnum.AGE_KICHI: OmikujiResult(
+        2,
+        ":poultry_leg: 揚げ吉 鳩を揚げると良いことあるよ!!"
+    ),
+    OmikujiEnum.KYO: OmikujiResult(
+        12,
+        ":cry: 凶 ちょっと慎重にいったほうがいいかも……"
+    ),
+    OmikujiEnum.DAI_KYO: OmikujiResult(
+        2,
+        ":crying_cat_face: 大凶 そういう時もあります……猫になって耐えましょう"
+    ),
+})
+
+
 def omikuji(client: BaseClient):
     """
     おみくじ結果を返す
     """
-    client.post(message=omikuji_draw())
+
+    logger.debug("%s called 'hato yoshiyoshi'", client.get_send_user())
+    logger.debug("%s app called 'hato yoshiyoshi'", client.get_type())
+    client.post(omikuji_draw(omikuji_results)[1].message)

tests/library/test_omikuji.py

@@ -4,7 +4,8 @@
 
 import unittest
 
-from library.omikuji import OmikujiResult, Omikuji
+from enum import Enum, auto
+from library.omikuji import OmikujiResult, OmikujiResults, draw
 
 
 class TestOmikuji(unittest.TestCase):
@@ -16,9 +17,19 @@ def test_omikuji_minimum(self):
         """
         おみくじの実装が正常か
         """
-        dummy_omikuji = Omikuji(entries={
-            'KICHI': OmikujiResult(1, "吉"),
-            'SUE_KICHI': OmikujiResult(1, "末吉"),
+        class DummyOmikujiEnum(Enum):
+            """
+            おみくじ設定(ダミー)
+            """
+            KICHI = auto()
+            SUE_KICHI = auto()
+
+        dummy_omikuji_results = OmikujiResults({
+            DummyOmikujiEnum.KICHI: OmikujiResult(1, "吉"),
+            DummyOmikujiEnum.SUE_KICHI: OmikujiResult(1, "末吉"),
         })
 
-        self.assertIn(dummy_omikuji.draw()[0], dummy_omikuji.entries.keys())
+        self.assertIn(
+            draw(dummy_omikuji_results)[0],
+            dummy_omikuji_results.keys()
+        )

tests/plugins/test_hato.py

@@ -10,12 +10,10 @@
 import requests_mock
 
 import slackbot_settings as conf
-from plugins.hato import split_command, amesh, altitude, yoshiyoshi, omikuji
+from plugins.hato import split_command, amesh, altitude, yoshiyoshi, omikuji, omikuji_results
 from tests.library.test_geo import set_mock
 from tests.plugins import TestClient
 
-from library.omikuji import omikuji as omikuji_instance
-
 
 class TestSplitCommand(unittest.TestCase):
     """
@@ -241,16 +239,14 @@ class TestOmikuji(unittest.TestCase):
 
     def test_omikuji(self):
         """
-        登録されたメッセージのいずれかが返ってくることのみ確認
+        設定したおみくじ結果のうち1つが返ってくる
         """
+
         client1 = TestClient()
         omikuji(client1)
         self.assertIn(
             client1.get_post_message(),
-            list(map(
-                lambda e: e.message,
-                omikuji_instance.entries.values()
-            ))
+            map(lambda e: e.message, omikuji_results.values())
         )