diff --git a/main.py b/main.py index ad98b41..203d3dc 100644 --- a/main.py +++ b/main.py @@ -12,6 +12,9 @@ import ffmpeg from dhooks import Webhook from fastapi import FastAPI, File, UploadFile from fastapi.responses import HTMLResponse +from pygments import highlight +from pygments.formatters import HtmlFormatter +from pygments.lexers import guess_lexer from settings import Settings @@ -71,13 +74,52 @@ def video_file_uploaded(file: UploadFile) -> Dict[str, str]: file_location = f"{Settings.upload_folder}/video/{file.filename}" height, width = find_video_resolution(file_location) screenshot_url = make_thumbnail_from_video(file_location, file.filename) - html_url = generate_html( + html_url = generate_html_for_videos( filename=file.filename, url=file_url, width=width, height=height, screenshot=screenshot_url ) hook.send(f"{Settings.domain}/{file.filename} was uploaded.") return {"html_url": f"{html_url}"} +def text_file_uploaded(file: UploadFile) -> Dict[str, str]: + """Save file to disk and return URL. + + Args: + file (UploadFile): Our file object. + + Returns: + Dict[str, str]: Returns URL for file. + """ + # Create folder if it doesn't exist. + Path(f"{Settings.upload_folder}/text").mkdir(parents=True, exist_ok=True) + + save_location = f"{Settings.upload_folder}/text/{file.filename}" + # Save file to disk. + with open(save_location, "wb+") as file_object: + file_object.write(file.file.read()) + + with open(save_location, encoding="utf-8") as file_object: + lines = file_object.read() + colored_text = highlight( + lines, + guess_lexer(lines), # Guess + HtmlFormatter( + style="fruity", # Dark style + linenos="table", # Output line numbers as a table w/ two cells, one with line numbers, other with code + full=True, # Use inline styles instead of CSS classes. + filename=f"{file.filename}", + ), + ) + + with open(f"{Settings.upload_folder}/{file.filename}.html", "w", encoding="utf-8") as file_object: + file_object.write(colored_text) + + html_url = f"{Settings.domain}/{file.filename}.html" + + hook.send(f"{html_url} was uploaded.") + return {"html_url": f"{html_url}"} + + @app.post("/uploadfiles/") async def upload_file(file: UploadFile = File(...)) -> Dict[str, str]: """Page for uploading files. @@ -92,11 +134,23 @@ async def upload_file(file: UploadFile = File(...)) -> Dict[str, str]: Returns: Dict[str, str]: Returns a dict with the filename or a link to the .html if it was a video. """ - # TODO: Add syntax highlighting for text. try: if file.content_type.startswith("video/"): return video_file_uploaded(file) - return normal_file_uploaded(file) + if ( + # TODO: This needs to be better and include more things. + file.content_type.startswith("text/") + or file.content_type.startswith("application/json") + or file.content_type.startswith("application/x-sh") + or file.content_type.startswith("application/xml") + ): + return text_file_uploaded(file) + + with open(f"{Settings.upload_folder}/{file.filename}", "wb+") as file_object: + file_object.write(file.file.read()) + + hook.send(f"{Settings.domain}/{file.filename} was uploaded.") + return {"html_url": f"{Settings.domain}/{file.filename}"} except Exception as e: # TODO: Change response code to 400. hook.send(f"Something went wrong for {Settings.domain}/{file.filename}:\n{e}") @@ -131,7 +185,7 @@ async def main(): """ -def generate_html(url: str, width: int, height: int, screenshot: str, filename: str) -> str: +def generate_html_for_videos(url: str, width: int, height: int, screenshot: str, filename: str) -> str: """Generate HTML for video files. This is what we will send to other people on Discord. diff --git a/poetry.lock b/poetry.lock index 8322ca8..8cb9b28 100644 --- a/poetry.lock +++ b/poetry.lock @@ -644,6 +644,14 @@ category = "dev" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +[[package]] +name = "pygments" +version = "2.10.0" +description = "Pygments is a syntax highlighting package written in Python." +category = "main" +optional = false +python-versions = ">=3.5" + [[package]] name = "pylint" version = "2.11.1" @@ -1000,7 +1008,7 @@ multidict = ">=4.0" [metadata] lock-version = "1.1" python-versions = "^3.9" -content-hash = "38172f23e6a3535e8146b631a26fa6b170b97b2bb5d1b88055148e4b016b000c" +content-hash = "113d57f8d58df9d307507a2200227a95e0784f4c78b9fb5f52fa67b0e1be222d" [metadata.files] aiohttp = [ @@ -1588,6 +1596,10 @@ pyflakes = [ {file = "pyflakes-2.3.1-py2.py3-none-any.whl", hash = "sha256:7893783d01b8a89811dd72d7dfd4d84ff098e5eed95cfa8905b22bbffe52efc3"}, {file = "pyflakes-2.3.1.tar.gz", hash = "sha256:f5bc8ecabc05bb9d291eb5203d6810b49040f6ff446a756326104746cc00c1db"}, ] +pygments = [ + {file = "Pygments-2.10.0-py3-none-any.whl", hash = "sha256:b8e67fe6af78f492b3c4b3e2970c0624cbf08beb1e493b2c99b9fa1b67a20380"}, + {file = "Pygments-2.10.0.tar.gz", hash = "sha256:f398865f7eb6874156579fdf36bc840a03cab64d1cde9e93d68f46a425ec52c6"}, +] pylint = [ {file = "pylint-2.11.1-py3-none-any.whl", hash = "sha256:0f358e221c45cbd4dad2a1e4b883e75d28acdcccd29d40c76eb72b307269b126"}, {file = "pylint-2.11.1.tar.gz", hash = "sha256:2c9843fff1a88ca0ad98a256806c82c5a8f86086e7ccbdb93297d86c3f90c436"}, diff --git a/pyproject.toml b/pyproject.toml index 5a48f94..84e945f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ gunicorn = "^20.1.0" uvicorn = "^0.15.0" python-dotenv = "^0.19.2" dhooks = {git = "https://github.com/kyb3r/dhooks.git"} +Pygments = "^2.10.0" [tool.poetry.dev-dependencies] pytest = "^5.2"