Update app.py

app.py

@@ -7,48 +7,43 @@ import os
 from PIL import Image
 import time
 import base64
-import validators  # URLバリデーション用
 
 
-def take_screenshot(url_or_html, wait_time=3, screenshot_type="url"):
+def take_screenshot_from_html(html_code, wait_time=3):
     """
-    URLまたはHTMLコードからフルページスクリーンショットを取得する関数。
+    HTMLコードからフルページスクリーンショットを取得する関数。
+    キャプチャ前に高さを3%拡張し、スクロールバーの有無に応じて調整。
 
     Args:
-        url_or_html (str): スクリーンショットを取得するURLまたはHTMLコード。
+        html_code (str): スクリーンショットを取得するHTMLコード。
         wait_time (int): ページ読み込みの待機時間（秒）。
-        screenshot_type (str): "url" または "html" を指定。
 
     Returns:
-        tuple: (PIL.Image.Image, str) スクリーンショット画像とメッセージのタプル。
+        tuple: (PIL.Image.Image, str) スクリーンショット画像とメッセージ。
                エラーの場合は (None, str) を返す。
     """
     options = Options()
-    options.add_argument("--headless=new")  # 新しいヘッドレスモード
+    options.add_argument("--headless=new")
     options.add_argument("--no-sandbox")
     options.add_argument("--disable-dev-shm-usage")
     options.add_argument("--disable-gpu")
-    options.add_argument("--window-position=-2400,-2400") # Windowsでのバグ回避
+    options.add_argument("--window-position=-2400,-2400")  # Windowsでのバグ回避
+
+    driver = None  # driverをtryの外で初期化
 
     try:
         driver = webdriver.Chrome(options=options)
+        html_base64 = base64.b64encode(html_code.encode('utf-8')).decode('utf-8')
+        driver.get(f"data:text/html;base64,{html_base64}")
 
-        if screenshot_type == "url":
-            # URLバリデーション
-            if not validators.url(url_or_html):
-                raise ValueError("無効なURLです。")
-            driver.get(url_or_html)
-
-        elif screenshot_type == "html":
-            # HTMLをbase64にエンコードしてdata URLとして読み込む
-            html_base64 = base64.b64encode(url_or_html.encode('utf-8')).decode('utf-8')
-            driver.get(f"data:text/html;base64,{html_base64}")
-        else:
-            raise ValueError("screenshot_typeは'url'または'html'を指定してください。")
+        time.sleep(wait_time)
 
-        time.sleep(wait_time)  # ページの読み込みを待機
+        # スクロールバーの有無をチェック
+        has_scrollbar = driver.execute_script(
+            "return document.documentElement.scrollHeight > document.documentElement.clientHeight;"
+        )
 
-        # フルページサイズを取得 (JavaScript)
+        # ページサイズを取得
         width = driver.execute_script(
             "return Math.max(document.body.scrollWidth, document.body.offsetWidth, "
             "document.documentElement.clientWidth, document.documentElement.scrollWidth, "
@@ -60,59 +55,47 @@ def take_screenshot(url_or_html, wait_time=3, screenshot_type="url"):
             "document.documentElement.offsetHeight);"
         )
 
-        # ウィンドウサイズを設定（巨大なページに対する制限付き）
-        driver.set_window_size(min(width + 100, 1920), min(height + 100, 5000))  # 最大幅1920px, 最大高さ5000px
-        time.sleep(1) # レンダリング待機
+        # 高さを3%拡張（スクロールバーがあればさらに調整）
+        extra_height = int(height * 0.03)
+        if has_scrollbar:
+            final_height = min(height + extra_height + 20, 5000)  # スクロールバー分を考慮、最大5000px
+        else:
+             final_height = min(height + extra_height, 5000)  # スクロールバーがなければ3%のみ、最大5000px
+
+        driver.set_window_size(min(width + 100, 1920), final_height)  # 最大幅1920
+        time.sleep(1)
 
-        # スクリーンショットを一時ファイルに保存
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp:
             temp_filename = temp.name
         driver.save_screenshot(temp_filename)
-
-        driver.quit()
-
-        # 画像を読み込んで返す
         image = Image.open(temp_filename)
-        os.unlink(temp_filename)  # 一時ファイルを削除
-
+        os.unlink(temp_filename)
         return image, "スクリーンショットを取得しました。"
 
-    except (WebDriverException, TimeoutException, ValueError) as e:
-        driver.quit()
+    except (WebDriverException, TimeoutException) as e:
         return None, f"エラー: {str(e)}"
     except Exception as e:
-        driver.quit()
         return None, f"予期しないエラー: {str(e)}"
-
-# Gradioインターフェース (URLとHTMLの両方に対応)
-with gr.Blocks(title="Web Screenshot Tool") as demo:
-    gr.Markdown("# ウェブページ & HTML スクリーンショットツール")
-    gr.Markdown("URLまたはHTMLコードを入力して、フルページのスクリーンショットを取得します。")
-
-    with gr.Tab("URLから取得"):
-        url_input = gr.Textbox(label="URL", placeholder="https://example.com")
-        url_wait_time = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="待機時間 (秒)")
-        url_button = gr.Button("スクリーンショットを取得 (URL)")
-        url_image_out = gr.Image(label="スクリーンショット (URL)")
-        url_text_out = gr.Textbox(label="結果 (URL)")
-
-        url_button.click(
-            take_screenshot,
-            inputs=[url_input, url_wait_time, gr.State("url")],  # "url"を指定
-            outputs=[url_image_out, url_text_out]
-        )
-
-    with gr.Tab("HTMLから取得"):
-        html_input = gr.Textbox(label="HTMLコード", placeholder="<p>Hello, world!</p>", lines=5)
-        html_wait_time = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="待機時間 (秒)")
-        html_button = gr.Button("スクリーンショットを取得 (HTML)")
-        html_image_out = gr.Image(label="スクリーンショット (HTML)")
-        html_text_out = gr.Textbox(label="結果 (HTML)")
-
-        html_button.click(
-            take_screenshot,
-            inputs=[html_input, html_wait_time, gr.State("html")],  # "html"を指定
-            outputs=[html_image_out, html_text_out]
-        )
+    finally:
+        if driver:  # driverが存在する場合のみquit()を実行
+            driver.quit()
+
+
+# Gradioインターフェース (HTML専用)
+with gr.Blocks(title="HTML Screenshot Tool") as demo:
+    gr.Markdown("# HTML スクリーンショットツール")
+    gr.Markdown("HTMLコードを入力して、レンダリングされたページのフルスクリーンショットを取得します。")
+
+    html_input = gr.Textbox(label="HTMLコード", placeholder="<p>Hello, world!</p>", lines=5)
+    wait_time_slider = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="待機時間 (秒)")
+    screenshot_button = gr.Button("スクリーンショットを取得")
+    image_output = gr.Image(label="スクリーンショット")
+    text_output = gr.Textbox(label="結果")
+
+    screenshot_button.click(
+        take_screenshot_from_html,
+        inputs=[html_input, wait_time_slider],
+        outputs=[image_output, text_output]
+    )
 
 demo.launch()