niming_igapi/utils/fileProcessor.py
2024-12-17 16:35:17 +00:00

108 lines
3.0 KiB
Python

import time
import os
import io
from typing import Tuple
import subprocess
from hashlib import sha512
from PIL import Image
from pillow_heif import register_heif_opener
import ffmpeg
from utils.const import FILE_MINE_TYPE, TMP_DIR
from utils.tbProcessor import easyExceptionHandler
register_heif_opener()
def image_conventer(filename:str, binary: bytes) -> int:
try:
fio = io.BytesIO(binary)
img:Image.Image = Image.open(fio)
img = img.convert("RGB")
img.save(filename, "JPEG", quality=95)
return 0
except Exception as e:
easyExceptionHandler(e)
return 1
def read_output(pipe, q):
""" 用於非阻塞讀取 ffmpeg 的 stdout """
while True:
data = pipe.read(4096)
if not data:
break
q.put(data)
q.put(None) # 標記輸出結束
def video_conventor(filename:str, oriFormat:str, binary:bytes) -> int:
try:
tmpfile = filename+"_tmp"
# write to tempfile
with open(tmpfile, "wb") as f:
f.write(binary)
# ffmpeg process
process:subprocess.Popen = (
ffmpeg
.input(tmpfile, format=oriFormat)
.output(filename, format='mp4')
.run_async(pipe_stdin=True, pipe_stdout=True, pipe_stderr=True)
)
process.wait()
# remove tempfile
os.remove(tmpfile)
return 0
except Exception as e:
easyExceptionHandler(e)
return 1
def file_writer(filename:str, binary:bytes):
with open(filename, "wb") as f:
f.write(binary)
def file_saver(ftype:str, binary:bytes) -> Tuple[str, int]:
"""
ftype -> minetype
binary -> file binary
"""
# 獲取副檔名
ext = None
for t in FILE_MINE_TYPE:
if t == ftype:
ext = FILE_MINE_TYPE[t]
if ext is None:
return "Invalid file type", 1
# 如果不是 IG 本身支援的檔案 -> 轉檔
filename = sha512( str(time.time()).encode() ).hexdigest() # 暫存檔名稱
opt = "" # output file name
if not ( ftype == "image/jpg" or ftype == "image/webp" or \
ftype == "video/mp4" ):
# 轉圖片
if ftype.startswith("image"):
opt = os.path.abspath(os.path.join(TMP_DIR, filename+".jpg"))
err = image_conventer(opt, binary)
if err: # 發生錯誤
return "File convert error", 1
# 轉影片
elif ftype.startswith("video"):
opt = os.path.abspath(os.path.join(TMP_DIR, filename+".mp4"))
err = video_conventor(opt, ext, binary)
if err:
return "File convert error", 1
# 轉檔完成
return opt, 0
else: # 如果是 IG 本身支援的檔案 -> 存檔
opt = os.path.abspath(os.path.join(TMP_DIR, filename+"."+ext))
file_writer(opt, binary)
return opt, 0