improve signature for the backend

streamlit_pdf_viewer/__init__.py

@@ -1,7 +1,9 @@
-import streamlit as st
-import streamlit.components.v1 as components
-import os
 import base64
+import os
+from pathlib import Path
+from typing import Union
+
+import streamlit.components.v1 as components
 
 _RELEASE = False
 
@@ -17,18 +19,21 @@
         "pdf_viewer", path=build_dir)
 
 
-def pdf_viewer(input: str, width="100%", height="700", key=None):
-    # "default" is a special argument that specifies the initial return
-    # value of the component before the user has interacted with it.
-    base64_pdf = base64.b64encode(input).decode('utf-8')
+def pdf_viewer(input: Union[str, Path, bytes], width="100%", height="700", key=None):
+    if type(input) is not bytes:
+        with open(input, 'rb') as fo:
+            binary = fo.read()
+    else:
+        binary = input
+
+    base64_pdf = base64.b64encode(binary).decode('utf-8')
     component_value = _component_func(binary=base64_pdf, width=width, height=height, key=key, default=0)
-    # We could modify the value returned from the component if we wanted.
-    # There's no need to do this in our simple example - but it's an option.
     return component_value
 
 
+# viewer = pdf_viewer("resources/test.pdf", height="700", width="700")
+
 with open("resources/test.pdf", 'rb') as fo:
     binary = fo.read()
 
-viewer = pdf_viewer(binary, height="700", width="700")
-st.write(viewer)
+viewer = pdf_viewer(binary, height="700", width="700")