Update app.py

app.py

@@ -11,8 +11,9 @@ from reportlab.lib.units import cm
 
 with open('pokemon.json', 'r') as f:
     pokemons = json.load(f)
-
+pokemons_types = ["모든 타입"] + sorted(set([t for poke in pokemons for t in poke['types']]))
 GEN_RANGE = {
+    "모든 세대": [1, 1017],
     "1세대": [1, 151],
     "2세대": [152, 251],
     "3세대": [252, 386],
@@ -25,14 +26,17 @@ GEN_RANGE = {
 }
 
 generation = gr.Dropdown(
-            [f"{k}세대" for k in range(1, 10)], value="1세대", label="포켓몬 세대", info="원하는 포켓몬 세대를 선택하세요."
+            [f"{k}세대" for k in range(1, 10)] + ["모든 세대"], value="1세대", label="포켓몬 세대", info="원하는 포켓몬 세대를 선택하세요."
+        )
+poke_types = gr.Dropdown(
+            pokemons_types, value="모든 타입", label="포켓몬 타입", info="원하는 포켓몬 타입을 선택하세요."
         )
 
 download = gr.File(label="파란 글씨를 눌러서 다운로드 받으세요.")
 text = gr.DataFrame(label="포켓몬 리스트")
 
-def write_pdf(gen, progress=gr.Progress()):
-    filename = f'포켓몬{gen}.pdf'
+def write_pdf(gen, poke_type, progress=gr.Progress()):
+    filename = f'포켓몬{gen}_{poke_type}.pdf'
     pdfmetrics.registerFont(TTFont("나눔고딕", "NanumGothic.ttf"))
     # PDF 문서를 생성합니다.
     pdfmetrics.registerFont(TTFont("나눔고딕", "NanumGothic.ttf"))
@@ -60,30 +64,35 @@ def write_pdf(gen, progress=gr.Progress()):
         types = pokemons[k-1]['types']
         image_path = pokemons[k-1]['image_path']
 
-        data_dict.append(
-            dict(이름=name, No=number, 타입='+'.join(types))
-        )
+        if poke_type == "모든 타입":
+            condition = True
+        else:
+            condition = poke_type in types
+        if condition:
+            data_dict.append(
+                dict(이름=name, No=number, 타입='+'.join(types))
+            )
 
-        # PDF에 이미지를 추가합니다. 이미지 파일의 경로를 지정하세요.
-        image = Image(image_path, width=13.5*cm, height=13.5*cm)
+            # PDF에 이미지를 추가합니다. 이미지 파일의 경로를 지정하세요.
+            image = Image(image_path, width=13.5*cm, height=13.5*cm)
 
-        # 테이블 데이터를 준비합니다.
-        data = [
-            [number],
-            [image],  # 2번째 행에 이미지를 추가할 것입니다.
-            ['+'.join(types)],
-            [name],
-        ]
-        # 테이블 객체를 생성합니다.
-        table = Table(data)
-        table.setStyle(style)
-        story.append(table)
+            # 테이블 데이터를 준비합니다.
+            data = [
+                [number],
+                [image],  # 2번째 행에 이미지를 추가할 것입니다.
+                ['+'.join(types)],
+                [name],
+            ]
+            # 테이블 객체를 생성합니다.
+            table = Table(data)
+            table.setStyle(style)
+            story.append(table)
 
     doc.build(story)
     df = pd.DataFrame(data_dict)
     return df[['No', '이름', '타입']], filename
 
-demo = gr.Interface(write_pdf, generation, [text, download], title="대치동 포켓몬 도감 생성기",
+demo = gr.Interface(write_pdf, [generation, poke_types], [text, download], title="대치동 포켓몬 도감 생성기",
                     description="원하는 포켓몬 세대를 선택하고, 다운로드를 눌러주세요.")
 demo.queue(concurrency_count=3)
 demo.launch()