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

dev-hato · Mar 9, 2022 · 55febf8 · 55febf8
1 parent ff6a85e
commit 55febf8
Show file tree

Hide file tree

Showing 4 changed files with 87 additions and 70 deletions.
diff --git a/library/omikuji.py b/library/omikuji.py
@@ -4,7 +4,8 @@
 おみくじを返す
 """
 
-from typing import Hashable, Tuple
+from enum import Enum, auto
+from typing import Tuple, TypeVar, Generic
 from dataclasses import dataclass
 from random import choices
 
@@ -26,66 +27,26 @@ 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')
 
-# 以下おみくじの設定
 
-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: 大凶 そういう時もあります……猫になって耐えましょう"
-    ),
+class OmikujiResults(dict[TOmikujiEnum, OmikujiResult]):
+    """
+    おみくじ結果と出やすさを管理する辞書
+    """
 
-})
+    def __init__(self, *args, **kwargs):
+        super(OmikujiResults, self).__init__(*args, **kwargs)
 
 
-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]
diff --git a/plugins/hato.py b/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 OmikujiResults, OmikujiResults, draw as omikuji_draw
 logger = getLogger(__name__)
 
 
@@ -277,8 +278,57 @@ def yoshiyoshi(client: BaseClient):
     client.post('よしよし')
 
 
+# 以下おみくじの設定
+# TODO: dev-hato/hato-bot#876
+class OmikujiEnum(Enum):
+    """
+    おみくじの結果一覧
+    """
+    DAI_KICHI = auto()
+    CHU_KICHI = auto()
+    SHO_KICHI = auto()
+    KICHI = auto()
+    HATO_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(
+        137,
+        ":smily: 小吉 なんとなくうまくいくかも!?"
+    ),
+    OmikujiEnum.HATO_KICHI: OmikujiResult(
+        2,
+        ":dove_of_peace: 鳩吉 お前が鳩になる番だ!!羽ばたけ!!!飛べ!!!!唐揚げになれ!!!!!"
+    ),
+    OmikujiEnum.KYO: OmikujiResult(
+        12,
+        ":cry: 凶 ちょっと慎重にいったほうがいいかも……"
+    ),
+    OmikujiEnum.DAI_KYO: OmikujiResult(
+        2,
+        ":crying_cat_face: 大凶 そういう時もあります……猫になって耐えましょう"
+    ),
+})
+
+
 def omikuji(client: BaseClient):
     """
     おみくじ結果を返す
     """
-    client.post(message=omikuji_draw())
+    client.post(omikuji_draw(
+        self.omikuji_results
+    )[1].message)
diff --git a/tests/library/test_omikuji.py b/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,16 @@ 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()
+        )
diff --git a/tests/plugins/test_hato.py b/tests/plugins/test_hato.py
@@ -14,7 +14,7 @@
 from tests.library.test_geo import set_mock
 from tests.plugins import TestClient
 
-from library.omikuji import omikuji as omikuji_instance
+from library.omikuji import omikuji as omikuji_results
 
 
 class TestSplitCommand(unittest.TestCase):
@@ -241,16 +241,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())
         )