Compare commits
5 Commits
v0.0.2
...
363b9a3d45
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
363b9a3d45 | ||
|
|
bc3a7500e2 | ||
|
|
8dd1f55fdd | ||
| b62473c468 | |||
| d8bf1818d7 |
1
.gitignore
vendored
1
.gitignore
vendored
@@ -1,4 +1,5 @@
|
|||||||
__pycache__/
|
__pycache__/
|
||||||
|
*.vscode
|
||||||
.venv/
|
.venv/
|
||||||
dist/
|
dist/
|
||||||
*.egg-info/
|
*.egg-info/
|
||||||
|
|||||||
14
README.md
14
README.md
@@ -16,24 +16,26 @@ python -m pip install git+https://git.tyukalov.su/nikita/mab
|
|||||||
|
|
||||||
## Basic usage
|
## Basic usage
|
||||||
|
|
||||||
This is the most simple bot you can create
|
This is the most simple bot you can create. It would respond to any message
|
||||||
|
that starts with `!test`, `!hello` or `!hi`.
|
||||||
|
|
||||||
```python
|
```python
|
||||||
import asyncio
|
import asyncio
|
||||||
from mab import MatrixBot, MatrixBotConfig
|
from mab import MatrixBot, MatrixBotConfig
|
||||||
|
from mab import TextCommandFilter
|
||||||
|
|
||||||
from nio import MatrixRoom, MatrixMessageText
|
from nio import MatrixRoom, MatrixMessageText
|
||||||
|
|
||||||
bot: MatrixBot
|
bot: MatrixBot
|
||||||
|
|
||||||
async def on_room_message_text(room: MatrixRoom, event: RoomMessageText) -> None:
|
async def on_valid_command(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||||
global bot
|
global bot
|
||||||
# do not respond to ourselves
|
# do not respond to ourselves
|
||||||
if event.sender == bot.get_client().user_id:
|
if event.sender == bot.get_client().user_id:
|
||||||
return
|
return
|
||||||
text = f"You message contains {len(event.body)} symbols"
|
text = f"Your message contains {len(event.body)} symbols"
|
||||||
await bot.send_text_to_room(room.room_id, text)
|
await bot.send_text_to_room(room.room_id, text)
|
||||||
|
|
||||||
|
|
||||||
async def main() -> None:
|
async def main() -> None:
|
||||||
global bot
|
global bot
|
||||||
# create and start the bot
|
# create and start the bot
|
||||||
@@ -43,6 +45,10 @@ async def main() -> None:
|
|||||||
storage_directory=Path("storage_nagibator666")
|
storage_directory=Path("storage_nagibator666")
|
||||||
)
|
)
|
||||||
bot = MatrixBot(matrix_bot_config)
|
bot = MatrixBot(matrix_bot_config)
|
||||||
|
bot.add_event_callback(
|
||||||
|
on_valid_command,
|
||||||
|
TextCommandFilter(["test", "hello", "hi"])
|
||||||
|
)
|
||||||
bot.start()
|
bot.start()
|
||||||
# wait for Ctrl+C
|
# wait for Ctrl+C
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|||||||
|
|
||||||
[project]
|
[project]
|
||||||
name = "mab"
|
name = "mab"
|
||||||
version = "0.0.2"
|
version = "0.1.0"
|
||||||
authors = [
|
authors = [
|
||||||
{ name = "Tyukalov Nikita", email = "nikita@tyukalov.su" }
|
{ name = "Tyukalov Nikita", email = "nikita@tyukalov.su" }
|
||||||
]
|
]
|
||||||
@@ -47,6 +47,8 @@ dependencies = [
|
|||||||
"unpaddedbase64==2.1.0",
|
"unpaddedbase64==2.1.0",
|
||||||
"vodozemac==0.10.0",
|
"vodozemac==0.10.0",
|
||||||
"yarl==1.24.5",
|
"yarl==1.24.5",
|
||||||
|
"python-magic==0.4.27",
|
||||||
|
"pillow==12.3.0"
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.urls]
|
[project.urls]
|
||||||
|
|||||||
@@ -5,6 +5,9 @@ from .types import MatrixBotConfig
|
|||||||
|
|
||||||
from .bot import MatrixBot
|
from .bot import MatrixBot
|
||||||
|
|
||||||
|
from .filters.base import *
|
||||||
|
from .filters.text import *
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
# module names
|
# module names
|
||||||
"bot",
|
"bot",
|
||||||
@@ -14,5 +17,16 @@ __all__ = [
|
|||||||
"MatrixBotConfig",
|
"MatrixBotConfig",
|
||||||
|
|
||||||
# .bot
|
# .bot
|
||||||
"MatrixBot"
|
"MatrixBot",
|
||||||
|
|
||||||
|
# .filters.base
|
||||||
|
"BaseEventFilter",
|
||||||
|
|
||||||
|
# .filters.text
|
||||||
|
"TextFilter",
|
||||||
|
"FormattedTextFilter",
|
||||||
|
"TextContainsFilter",
|
||||||
|
"TextStartsWithFilter",
|
||||||
|
"TextEndsWithFilter",
|
||||||
|
"TextCommandFilter",
|
||||||
]
|
]
|
||||||
449
src/mab/bot.py
449
src/mab/bot.py
@@ -1,5 +1,5 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import aiofiles
|
import aiofiles, aiofiles.os
|
||||||
import aioconsole
|
import aioconsole
|
||||||
import traceback
|
import traceback
|
||||||
import logging
|
import logging
|
||||||
@@ -7,23 +7,31 @@ import time
|
|||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
|
import io
|
||||||
|
import magic
|
||||||
from html.parser import HTMLParser
|
from html.parser import HTMLParser
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
from typing import Any, Callable, Awaitable
|
from typing import Any, Callable, Awaitable
|
||||||
|
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
from nio import AsyncClient, AsyncClientConfig, SyncResponse
|
from nio import AsyncClient, AsyncClientConfig, SyncResponse
|
||||||
from nio import LoginResponse, LoginError, WhoamiResponse, WhoamiError
|
from nio import LoginResponse, LoginError, WhoamiResponse, WhoamiError
|
||||||
|
|
||||||
from nio import RoomSendResponse, RoomSendError
|
from nio import RoomSendResponse, RoomSendError
|
||||||
|
from nio import UploadResponse, UploadError
|
||||||
|
|
||||||
from nio import OlmUnverifiedDeviceError
|
from nio import OlmUnverifiedDeviceError
|
||||||
|
|
||||||
from nio import MatrixInvitedRoom, InviteMemberEvent
|
from nio import MatrixInvitedRoom, InviteMemberEvent
|
||||||
from nio import JoinResponse
|
from nio import JoinResponse
|
||||||
|
|
||||||
import nio.events
|
from .filters.base import BaseEventFilter
|
||||||
|
|
||||||
|
from nio.events.room_events import Event as RoomEvemt
|
||||||
|
|
||||||
from .types import *
|
from .types import *
|
||||||
|
from .utils import Utils
|
||||||
|
|
||||||
|
|
||||||
class MatrixBot:
|
class MatrixBot:
|
||||||
@@ -106,31 +114,59 @@ class MatrixBot:
|
|||||||
return f"@{localpart}:{homeserver_name}"
|
return f"@{localpart}:{homeserver_name}"
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _check_and_remove_html(text_to_check: str) -> tuple[bool, str]:
|
def _process_html_text(text: str, is_html: bool | None) -> dict:
|
||||||
"""Checks if `text_to_check` is HTML and sanitizes it.
|
"""Processes `text` and `is_html` arguments of `send_*` functions and
|
||||||
|
prepares `content` keys related to text data inside the message.
|
||||||
Returns:
|
|
||||||
tuple[bool, str] where the first element is True if `text_to_check` contains
|
Args:
|
||||||
valid HTML, and the second element is text without HTML (or just copy of
|
text - `text` as passed to the `send_*` function
|
||||||
`text_to_check` if it does not contain HTML)
|
is_html - `is_html` as passed to the `send_*` function
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
`dict` that contains keys required for the text message to be
|
||||||
|
displayed correctly.
|
||||||
"""
|
"""
|
||||||
has_tags = False
|
# text w/ and w/o HTML
|
||||||
text_fragments = []
|
text_w_html: str | None = None
|
||||||
class Extractor(HTMLParser):
|
text_wo_html: str | None = None
|
||||||
def handle_starttag(self, tag, attrs):
|
# check if HTML is present
|
||||||
nonlocal has_tags
|
if is_html or is_html is None:
|
||||||
has_tags = True
|
# will be True if `text` contains HTML tags
|
||||||
def handle_data(self, data):
|
has_tags = False
|
||||||
text_fragments.append(data)
|
# text fragments without tags
|
||||||
parser = Extractor(convert_charrefs=True)
|
text_fragments = []
|
||||||
parser.feed(text_to_check)
|
# temporary extractor class
|
||||||
try:
|
class Extractor(HTMLParser):
|
||||||
if has_tags:
|
def handle_starttag(self, tag, attrs):
|
||||||
return (True, " ".join("".join(text_fragments).split()))
|
nonlocal has_tags
|
||||||
except:
|
has_tags = True
|
||||||
traceback.print_exc()
|
def handle_data(self, data):
|
||||||
return (False, text_to_check)
|
text_fragments.append(data)
|
||||||
|
# parse
|
||||||
|
parser = Extractor(convert_charrefs=True)
|
||||||
|
try:
|
||||||
|
parser.feed(text)
|
||||||
|
except:
|
||||||
|
traceback.print_exc()
|
||||||
|
text_wo_html = text
|
||||||
|
is_html = has_tags
|
||||||
|
if is_html:
|
||||||
|
text_wo_html = " ".join(text_fragments)
|
||||||
|
text_w_html = text
|
||||||
|
if not is_html:
|
||||||
|
text_wo_html = text
|
||||||
|
# return the content
|
||||||
|
if text_w_html:
|
||||||
|
return {
|
||||||
|
"body": text_wo_html,
|
||||||
|
"format": "org.matrix.custom.html",
|
||||||
|
"formatted_body": text_w_html
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
return {
|
||||||
|
"body": text_wo_html
|
||||||
|
}
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
async def _wait_for_task_and_stop_event(payload_task: asyncio.Task, stop_wait_task: asyncio.Task, print_exc: bool = False) -> Any:
|
async def _wait_for_task_and_stop_event(payload_task: asyncio.Task, stop_wait_task: asyncio.Task, print_exc: bool = False) -> Any:
|
||||||
"""Waits for `payload_task` to finish. Uses `stop_wait_task` to check if payload task
|
"""Waits for `payload_task` to finish. Uses `stop_wait_task` to check if payload task
|
||||||
@@ -176,7 +212,8 @@ class MatrixBot:
|
|||||||
print(f" * {k} = {kwargs[k]}")
|
print(f" * {k} = {kwargs[k]}")
|
||||||
except:
|
except:
|
||||||
traceback.print_exc()
|
traceback.print_exc()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
async def _write_next_batch(self, next_batch: str) -> None:
|
async def _write_next_batch(self, next_batch: str) -> None:
|
||||||
"""Writes `next_batch` value to disk."""
|
"""Writes `next_batch` value to disk."""
|
||||||
@@ -213,6 +250,154 @@ class MatrixBot:
|
|||||||
j = json.loads(await f.read())
|
j = json.loads(await f.read())
|
||||||
return j
|
return j
|
||||||
|
|
||||||
|
async def _upload_file(self, path: Path | str, mime_type: str, *, filename: str | None = None) -> tuple[UploadResponse, dict, int]:
|
||||||
|
"""Uploads the file with encryption enabled. UploadError will be raised
|
||||||
|
as RuntimeError.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
path - path to the file to upload
|
||||||
|
mime_type - mime-type to use
|
||||||
|
filename - filename to use (base name of the `path` will be used if not set)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
1. UploadResponse as returned by `_upload_binary_data(...)`
|
||||||
|
2. dict as returned by `_upload_binary_data(...)`
|
||||||
|
3. file size in bytes
|
||||||
|
"""
|
||||||
|
# get basic file info
|
||||||
|
file_info = await aiofiles.os.stat(path)
|
||||||
|
if not filename:
|
||||||
|
filename = os.path.basename(path)
|
||||||
|
# upload
|
||||||
|
async with aiofiles.open(path, "r+b") as f:
|
||||||
|
response, keys = await self._upload_provider(
|
||||||
|
f,
|
||||||
|
mime_type,
|
||||||
|
filename=filename,
|
||||||
|
filesize=file_info.st_size
|
||||||
|
)
|
||||||
|
return response, keys, file_info.st_size
|
||||||
|
|
||||||
|
async def _upload_provider(self, provider: Any, mime_type: str, *, filename: str, filesize: int) -> tuple[UploadResponse, dict]:
|
||||||
|
"""Uploads the file with encryption enabled. UploadError will be raised
|
||||||
|
as RuntimeError.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
provider - data provider as described in `AsyncClient.upload(...)` docs
|
||||||
|
mime_type - mime-type to use
|
||||||
|
filename - filename to use
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
1. UploadResponse as returned by `AsyncClient.upload(...)`
|
||||||
|
2. dict as returned by `AsyncClient.upload(...)`
|
||||||
|
"""
|
||||||
|
# upload
|
||||||
|
response, keys = await self._client.upload(
|
||||||
|
provider,
|
||||||
|
content_type=mime_type,
|
||||||
|
filename=filename,
|
||||||
|
encrypt=True,
|
||||||
|
monitor=None,
|
||||||
|
filesize=filesize
|
||||||
|
)
|
||||||
|
if isinstance(response, UploadError):
|
||||||
|
raise RuntimeError(response.message, response)
|
||||||
|
if not isinstance(response, UploadResponse):
|
||||||
|
raise RuntimeError(
|
||||||
|
"Upload result must be `UploadRespond`, but it's not",
|
||||||
|
response
|
||||||
|
)
|
||||||
|
return response, keys # type: ignore
|
||||||
|
|
||||||
|
async def _upload_thumbnail(self, thumbnail: Path | str | bytes | None) -> dict:
|
||||||
|
"""Upload thumbnail and return fields required to use it for video.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
thumbnail - path to the thumbnail file or raw `bytes` of the image.
|
||||||
|
You can actually pass `None`, empty dict will be returned
|
||||||
|
in this case.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
`dict` you should embed into your `content["info"]`
|
||||||
|
"""
|
||||||
|
if thumbnail is None:
|
||||||
|
return {}
|
||||||
|
# raw data
|
||||||
|
raw_bytes_io: io.BytesIO | None = None
|
||||||
|
if type(thumbnail) is bytes:
|
||||||
|
thumb_mime = magic.from_buffer(thumbnail, mime=True)
|
||||||
|
raw_bytes_io = io.BytesIO(thumbnail)
|
||||||
|
# file path
|
||||||
|
else:
|
||||||
|
thumb_mime = magic.from_file(thumbnail, mime=True)
|
||||||
|
# check if invalid MIME
|
||||||
|
if not thumb_mime.startswith("image/"):
|
||||||
|
raise RuntimeError(f"Thumbnail MIME is not image: `{thumb_mime}`")
|
||||||
|
# get thumbnail size
|
||||||
|
with Image.open(raw_bytes_io or thumbnail) as f:
|
||||||
|
width, height = f.size
|
||||||
|
# seek BytesIO to zero so we can reuse it
|
||||||
|
if raw_bytes_io is not None:
|
||||||
|
raw_bytes_io.seek(0)
|
||||||
|
# upload
|
||||||
|
if type(thumbnail) is bytes:
|
||||||
|
thumb_size = len(thumbnail)
|
||||||
|
thumb_resp, thumb_keys = await self._upload_provider(
|
||||||
|
raw_bytes_io,
|
||||||
|
thumb_mime,
|
||||||
|
filename=f"preview.{thumb_mime.split("/")[-1]}",
|
||||||
|
filesize=thumb_size
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
thumb_resp, thumb_keys, thumb_size = await self._upload_file(
|
||||||
|
thumbnail, # type: ignore
|
||||||
|
mime_type=thumb_mime,
|
||||||
|
filename=os.path.basename(str(thumbnail))
|
||||||
|
)
|
||||||
|
# return the result
|
||||||
|
return {
|
||||||
|
"thumbnail_file": {
|
||||||
|
"url": thumb_resp.content_uri,
|
||||||
|
"mimetype": thumb_mime,
|
||||||
|
**thumb_keys
|
||||||
|
},
|
||||||
|
"thumbnail_info": {
|
||||||
|
"w": width,
|
||||||
|
"h": height,
|
||||||
|
"size": thumb_size,
|
||||||
|
"mimetype": thumb_mime
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async def _send_content_to_room(self, room_id: str, content: dict) -> RoomSendResponse:
|
||||||
|
"""Sends `m.room.message` event to the room. Returns RoomSendResponse
|
||||||
|
as returned by `AsyncClient.room_send(...)`. Automatically verifies
|
||||||
|
all devices if this feature is enabled. Return values other than
|
||||||
|
`RoomSendResponse` are raised as `RuntimeError`."""
|
||||||
|
# try to send the message
|
||||||
|
try:
|
||||||
|
result = await self._client.room_send(
|
||||||
|
room_id=room_id,
|
||||||
|
message_type="m.room.message",
|
||||||
|
content=content
|
||||||
|
)
|
||||||
|
except OlmUnverifiedDeviceError:
|
||||||
|
if self._config.auto_verify_all_known_devices:
|
||||||
|
if not self.verify_all_known_devices():
|
||||||
|
raise
|
||||||
|
return await self._send_content_to_room(room_id, content)
|
||||||
|
else:
|
||||||
|
raise
|
||||||
|
# success
|
||||||
|
if type(result) is RoomSendResponse:
|
||||||
|
return result
|
||||||
|
# error
|
||||||
|
elif type(result) is RoomSendError:
|
||||||
|
raise RuntimeError(result)
|
||||||
|
# unknown error
|
||||||
|
else:
|
||||||
|
raise RuntimeError("Unknown error has occured", result)
|
||||||
|
|
||||||
#
|
#
|
||||||
# CALLBACKS
|
# CALLBACKS
|
||||||
#
|
#
|
||||||
@@ -244,13 +429,34 @@ class MatrixBot:
|
|||||||
except:
|
except:
|
||||||
self._logger.error(traceback.format_exc())
|
self._logger.error(traceback.format_exc())
|
||||||
|
|
||||||
|
async def _callback_filter_router(self, *args, **kwargs):
|
||||||
|
if len(args) != 2:
|
||||||
|
self._logger.debug("Can't process the event, not enough positional args")
|
||||||
|
await self._debug_event_callback(*args, **kwargs)
|
||||||
|
return
|
||||||
|
room = args[0]
|
||||||
|
event = args[1]
|
||||||
|
for filter in self._filters:
|
||||||
|
filter_object = filter[0]
|
||||||
|
filter_callback = filter[1]
|
||||||
|
filter_stop_after_this = filter[2]
|
||||||
|
if filter_object(room, event):
|
||||||
|
self._logger.debug(f"Filter {repr(filter_object)} matched")
|
||||||
|
try:
|
||||||
|
await filter_callback(room, event)
|
||||||
|
except:
|
||||||
|
self._logger.error(traceback.format_exc())
|
||||||
|
if filter_stop_after_this:
|
||||||
|
self._logger.debug(f"Filter {repr(filter_object)} stops matching")
|
||||||
|
break
|
||||||
|
|
||||||
#
|
#
|
||||||
# LIFECYCLE
|
# LIFECYCLE
|
||||||
#
|
#
|
||||||
def _setup_client_callbacks(self) -> None:
|
def _setup_client_callbacks(self) -> None:
|
||||||
"""Setup internal client callbacks"""
|
"""Setup internal client callbacks"""
|
||||||
self._client.add_response_callback(self._callback_sync, SyncResponse) # type: ignore
|
self._client.add_response_callback(self._callback_sync, SyncResponse) # type: ignore
|
||||||
|
self._client.add_event_callback(self._callback_filter_router, RoomEvemt) # type: ignore
|
||||||
if self._config.auto_join_any_room_on_invite:
|
if self._config.auto_join_any_room_on_invite:
|
||||||
self._client.add_event_callback(self._callback_autojoin, InviteMemberEvent) # type: ignore
|
self._client.add_event_callback(self._callback_autojoin, InviteMemberEvent) # type: ignore
|
||||||
|
|
||||||
@@ -389,6 +595,8 @@ class MatrixBot:
|
|||||||
#
|
#
|
||||||
def __init__(self, config: MatrixBotConfig) -> None:
|
def __init__(self, config: MatrixBotConfig) -> None:
|
||||||
# check if config is valid
|
# check if config is valid
|
||||||
|
if type(config.storage_directory) is str:
|
||||||
|
config.storage_directory = Path(config.storage_directory)
|
||||||
self._validate_bot_config(config) # may raise an Exception
|
self._validate_bot_config(config) # may raise an Exception
|
||||||
# save the config
|
# save the config
|
||||||
self._config: MatrixBotConfig = config
|
self._config: MatrixBotConfig = config
|
||||||
@@ -404,6 +612,7 @@ class MatrixBot:
|
|||||||
self._last_next_batch_dump: float = 0.0
|
self._last_next_batch_dump: float = 0.0
|
||||||
self._last_next_batch: str | None = None
|
self._last_next_batch: str | None = None
|
||||||
self._cb_password = self._default_password_callback
|
self._cb_password = self._default_password_callback
|
||||||
|
self._filters = []
|
||||||
|
|
||||||
def start(self) -> None:
|
def start(self) -> None:
|
||||||
"""Start the bot.
|
"""Start the bot.
|
||||||
@@ -451,18 +660,17 @@ class MatrixBot:
|
|||||||
result = True
|
result = True
|
||||||
return result
|
return result
|
||||||
|
|
||||||
def add_event_callback(self, callback: Callable[[Any], Awaitable[None]] | None, event_class: nio.events.Event) -> None:
|
def add_event_callback(self, callback: Callable[..., Awaitable[None]] | None, filter: BaseEventFilter, stop_after_this: bool = True) -> None:
|
||||||
"""Added event callback for events of specified class.
|
"""Add event callback for events that pass the filter."""
|
||||||
Use `None` instead of callback to print parameter types you need to use in your callback."""
|
|
||||||
if callback is None:
|
if callback is None:
|
||||||
callback = self._debug_event_callback
|
callback = self._debug_event_callback
|
||||||
self._client.add_event_callback(callback, event_class) # type: ignore
|
self._filters.append((filter, callback, stop_after_this))
|
||||||
|
|
||||||
def get_client(self) -> AsyncClient:
|
def get_client(self) -> AsyncClient:
|
||||||
"""Get AsyncClient in use"""
|
"""Get AsyncClient in use"""
|
||||||
return self._client
|
return self._client
|
||||||
|
|
||||||
async def send_text_to_room(self, room_id: str, text: str, is_html: bool | None = None, **kwargs) -> str:
|
async def send_text_to_room(self, room_id: str, text: str, is_html: bool | None = None) -> str:
|
||||||
"""Sends a text message to the room and handle HTML as specified.
|
"""Sends a text message to the room and handle HTML as specified.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
@@ -473,42 +681,147 @@ class MatrixBot:
|
|||||||
Returns:
|
Returns:
|
||||||
event_id of the message on success. Raises an exception on error.
|
event_id of the message on success. Raises an exception on error.
|
||||||
"""
|
"""
|
||||||
original_text = text
|
text_data = self._process_html_text(text, is_html)
|
||||||
# guess if text is HTML
|
|
||||||
if is_html is None:
|
|
||||||
is_html, text = self._check_and_remove_html(text)
|
|
||||||
# text is HTML
|
|
||||||
elif is_html:
|
|
||||||
_, text = self._check_and_remove_html(text)
|
|
||||||
# create `content` for `room_send()`
|
|
||||||
content = {
|
content = {
|
||||||
"msgtype": "m.text",
|
"msgtype": "m.text",
|
||||||
"body": text,
|
**text_data
|
||||||
**kwargs
|
|
||||||
}
|
}
|
||||||
if is_html:
|
return (await self._send_content_to_room(room_id, content)).event_id
|
||||||
content["format"] = "org.matrix.custom.html"
|
|
||||||
content["formatted_body"] = original_text
|
async def send_image_to_room(self,
|
||||||
# try to send the message
|
room_id: str,
|
||||||
try:
|
path: Path | str,
|
||||||
result = await self._client.room_send(
|
caption: str | None = None,
|
||||||
room_id=room_id,
|
is_html: bool | None = None) -> str:
|
||||||
message_type="m.room.message",
|
"""Send an image to the room.
|
||||||
content=content
|
|
||||||
|
This method sends specified image file to the specified room. It
|
||||||
|
will use mime-type deduced using `magic` module. Image dimension is
|
||||||
|
deduced using Pillow. The file is read asynchronously, so this
|
||||||
|
method won't block your event loop.
|
||||||
|
|
||||||
|
Please note that HTML is incorrectly interpreted for images in
|
||||||
|
some clients (as of September 6th, 2026).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
room_id - room_id to send the image to
|
||||||
|
path - path to the image file
|
||||||
|
caption - caption to use (will be hidden if it is the same as file basename)
|
||||||
|
is_html - whether the caption is HTML (autodetected if None)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
event_id
|
||||||
|
"""
|
||||||
|
# caption must not actually be empty
|
||||||
|
if caption is None or not caption.strip():
|
||||||
|
caption = os.path.basename(path)
|
||||||
|
is_html = False
|
||||||
|
# check if the file is image
|
||||||
|
mime_type: str = magic.from_file(path, mime=True)
|
||||||
|
if not mime_type.startswith("image/"):
|
||||||
|
raise RuntimeError(f"File {path} has non-image mime-type")
|
||||||
|
# get image size
|
||||||
|
with Image.open(path) as image:
|
||||||
|
width, height = image.size
|
||||||
|
# upload
|
||||||
|
response, keys, filesize = await self._upload_file(path, mime_type)
|
||||||
|
# prepare the content
|
||||||
|
content = {
|
||||||
|
"msgtype": "m.image",
|
||||||
|
"filename": os.path.basename(path),
|
||||||
|
**self._process_html_text(caption, is_html),
|
||||||
|
"file": {
|
||||||
|
"url": response.content_uri,
|
||||||
|
"mimetype": mime_type,
|
||||||
|
**keys
|
||||||
|
},
|
||||||
|
"info": {
|
||||||
|
"mimetype": mime_type,
|
||||||
|
"size": filesize,
|
||||||
|
"w": width,
|
||||||
|
"h": height
|
||||||
|
}
|
||||||
|
}
|
||||||
|
# send
|
||||||
|
return (await self._send_content_to_room(room_id, content)).event_id
|
||||||
|
|
||||||
|
async def send_video_to_room(self,
|
||||||
|
room_id: str,
|
||||||
|
path: Path | str,
|
||||||
|
props: VideoFileProperties | None = None,
|
||||||
|
caption: str | None = None,
|
||||||
|
is_html: bool | None = None) -> str:
|
||||||
|
"""Send a video to the room.
|
||||||
|
|
||||||
|
This method sends specified video file to the specified room. It
|
||||||
|
will use mime-type deduced using `magic` module. Most video
|
||||||
|
properties must be passed using `props` argument.
|
||||||
|
|
||||||
|
If you pass `None` as `props` then system `ffprobe` will be called
|
||||||
|
to get video dimensions and duration. You can disable this behavior
|
||||||
|
by disabling `auto_ffprobe` in `MatrixBotConfig`.
|
||||||
|
|
||||||
|
If you pass `None` as `props.thumbnail_path` then thumbnail will be
|
||||||
|
generated automatically. You can disable this behavior by disabling
|
||||||
|
`auto_ffmpeg` in `MatrixBotConfig`.
|
||||||
|
|
||||||
|
Please note that HTML is incorrectly interpreted for videos in
|
||||||
|
some clients (as of September 6th, 2026).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
room_id - room_id to send the video to
|
||||||
|
path - path to the video file
|
||||||
|
props - video properties
|
||||||
|
caption - caption to use (will be hidden if it is the same as file basename)
|
||||||
|
is_html - whether the caption is HTML (autodetected if None)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
event_id
|
||||||
|
"""
|
||||||
|
# caption must not actually be empty
|
||||||
|
if caption is None or not caption.strip():
|
||||||
|
caption = os.path.basename(path)
|
||||||
|
is_html = False
|
||||||
|
# check if the file is video
|
||||||
|
mime_type: str = magic.from_file(path, mime=True)
|
||||||
|
if not mime_type.startswith("video/"):
|
||||||
|
raise RuntimeError(f"File {path} has non-video mime-type")
|
||||||
|
# properties must not be empty
|
||||||
|
if props is None:
|
||||||
|
if not self._config.auto_ffprobe:
|
||||||
|
raise RuntimeError("No `props` are supplied and `auto_ffprobe` is disabled")
|
||||||
|
props = await Utils.get_video_props_using_ffprobe(path)
|
||||||
|
# thumbnail data to use
|
||||||
|
thumbnail_data = {}
|
||||||
|
# autogenerate thumbnail if allowed
|
||||||
|
if not props.thumbnail and self._config.auto_ffmpeg:
|
||||||
|
props.thumbnail = await Utils.get_video_thumbnail(
|
||||||
|
path,
|
||||||
|
timestamp=props.duration_ms / 1000 / 2,
|
||||||
|
max_size=400
|
||||||
)
|
)
|
||||||
except OlmUnverifiedDeviceError:
|
# upload the thumbnail if specified
|
||||||
if self._config.auto_verify_all_known_devices:
|
thumbnail_data = await self._upload_thumbnail(props.thumbnail)
|
||||||
if not self.verify_all_known_devices():
|
# upload
|
||||||
raise
|
response, keys, filesize = await self._upload_file(path, mime_type)
|
||||||
return await self.send_text_to_room(room_id, original_text, is_html, **kwargs)
|
# prepare the content
|
||||||
else:
|
content = {
|
||||||
raise
|
"msgtype": "m.video",
|
||||||
# success
|
"filename": os.path.basename(path),
|
||||||
if type(result) is RoomSendResponse:
|
**self._process_html_text(caption, is_html),
|
||||||
return result.event_id
|
"file": {
|
||||||
# error
|
"url": response.content_uri,
|
||||||
elif type(result) is RoomSendError:
|
"mimetype": mime_type,
|
||||||
raise RuntimeError(result)
|
**keys
|
||||||
# unknown error
|
},
|
||||||
else:
|
"info": {
|
||||||
raise RuntimeError("Unknown error has occured", result)
|
"duration": props.duration_ms,
|
||||||
|
"mimetype": mime_type,
|
||||||
|
"size": filesize,
|
||||||
|
"w": props.width_px,
|
||||||
|
"h": props.height_px,
|
||||||
|
**thumbnail_data
|
||||||
|
}
|
||||||
|
}
|
||||||
|
# send
|
||||||
|
return (await self._send_content_to_room(room_id, content)).event_id
|
||||||
139
src/mab/filters/base.py
Normal file
139
src/mab/filters/base.py
Normal file
@@ -0,0 +1,139 @@
|
|||||||
|
from abc import ABC, abstractmethod
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from nio import MatrixRoom, Event
|
||||||
|
|
||||||
|
class BaseEventFilter(ABC):
|
||||||
|
"""Base class for all message filters"""
|
||||||
|
_logger = logging.Logger("EventFilter")
|
||||||
|
|
||||||
|
# AND
|
||||||
|
def __and__(self, other):
|
||||||
|
if not isinstance(other, BaseEventFilter):
|
||||||
|
raise TypeError(f"BaseEventFilter can't be ANDed against {type(other)}")
|
||||||
|
return CompoundEventFilter(
|
||||||
|
CompoundEventFilter.OPERATOR_AND,
|
||||||
|
[self, other]
|
||||||
|
)
|
||||||
|
|
||||||
|
def __rand__(self, other):
|
||||||
|
return self.__and__(other)
|
||||||
|
|
||||||
|
# OR
|
||||||
|
def __or__(self, other):
|
||||||
|
if not isinstance(other, BaseEventFilter):
|
||||||
|
raise TypeError(f"BaseEventFilter can't be ORed against {type(other)}")
|
||||||
|
return CompoundEventFilter(
|
||||||
|
CompoundEventFilter.OPERATOR_OR,
|
||||||
|
[self, other]
|
||||||
|
)
|
||||||
|
|
||||||
|
def __ror__(self, other):
|
||||||
|
return self.__ror__(other)
|
||||||
|
|
||||||
|
# XOR
|
||||||
|
def __xor__(self, other):
|
||||||
|
if not isinstance(other, BaseEventFilter):
|
||||||
|
raise TypeError(f"BaseEventFilter can't be XORed against {type(other)}")
|
||||||
|
return CompoundEventFilter(
|
||||||
|
CompoundEventFilter.OPERATOR_XOR,
|
||||||
|
[self, other]
|
||||||
|
)
|
||||||
|
|
||||||
|
def __rxor__(self, other):
|
||||||
|
return self.__xor__(other)
|
||||||
|
|
||||||
|
# INVERT
|
||||||
|
def __invert__(self):
|
||||||
|
return CompoundEventFilter(
|
||||||
|
CompoundEventFilter.OPERATOR_INVERT,
|
||||||
|
[self]
|
||||||
|
)
|
||||||
|
|
||||||
|
# PAYLOAD
|
||||||
|
@abstractmethod
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
"""This method must be redefined in derived classes to improve
|
||||||
|
debugging experience.
|
||||||
|
"""
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||||
|
"""This abstract method must be redefined in derived classes so that
|
||||||
|
the filter operates according to its description. This method must
|
||||||
|
not raise exceptions. In case of exception it should log it using
|
||||||
|
`self._logger` and return False
|
||||||
|
|
||||||
|
Args:
|
||||||
|
event - the event to check againts this filter
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if the event satisfies this filter
|
||||||
|
False if the event does not satisfy this filter
|
||||||
|
"""
|
||||||
|
pass
|
||||||
|
|
||||||
|
class CompoundEventFilter(BaseEventFilter):
|
||||||
|
"""Event filter that consists of multiple filters"""
|
||||||
|
OPERATOR_AND = "and"
|
||||||
|
OPERATOR_OR = "or"
|
||||||
|
OPERATOR_XOR = "xor"
|
||||||
|
OPERATOR_INVERT = "invert"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_operator_valid(op: str) -> bool:
|
||||||
|
"""Returns True if operator `op` is a valid operator"""
|
||||||
|
return op in [
|
||||||
|
CompoundEventFilter.OPERATOR_AND,
|
||||||
|
CompoundEventFilter.OPERATOR_OR,
|
||||||
|
CompoundEventFilter.OPERATOR_XOR,
|
||||||
|
CompoundEventFilter.OPERATOR_INVERT
|
||||||
|
]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_elements_count_valid(op: str, count: int) -> bool:
|
||||||
|
"""Returns True if operator `op` may take `count` arguments"""
|
||||||
|
return count in {
|
||||||
|
CompoundEventFilter.OPERATOR_AND: [2],
|
||||||
|
CompoundEventFilter.OPERATOR_OR: [2],
|
||||||
|
CompoundEventFilter.OPERATOR_XOR: [2],
|
||||||
|
CompoundEventFilter.OPERATOR_INVERT: [1],
|
||||||
|
}[op]
|
||||||
|
|
||||||
|
def __init__(self, operator: str, arguments: list[BaseEventFilter]):
|
||||||
|
super().__init__()
|
||||||
|
if not self._is_operator_valid(operator):
|
||||||
|
raise RuntimeError(f"Invalid operator `{operator}`")
|
||||||
|
if not self._is_elements_count_valid(operator, len(arguments)):
|
||||||
|
raise RuntimeError(f"Operator `{operator}` does not take `{len(arguments)}` arguments")
|
||||||
|
self._operator = operator
|
||||||
|
self._arguments = list(arguments)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
expression = "False"
|
||||||
|
reprs = [repr(a) for a in self._arguments]
|
||||||
|
if self._operator == CompoundEventFilter.OPERATOR_AND:
|
||||||
|
expression = " & ".join(reprs)
|
||||||
|
elif self._operator == CompoundEventFilter.OPERATOR_OR:
|
||||||
|
expression = " | ".join(reprs)
|
||||||
|
elif self._operator == CompoundEventFilter.OPERATOR_XOR:
|
||||||
|
expression = " ^ ".join(reprs)
|
||||||
|
elif self._operator == CompoundEventFilter.OPERATOR_INVERT:
|
||||||
|
expression = f"~{reprs[0]}"
|
||||||
|
return f"({expression})"
|
||||||
|
|
||||||
|
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||||
|
evaluated = [arg(room, event) for arg in self._arguments]
|
||||||
|
if self._operator == self.OPERATOR_AND:
|
||||||
|
return all(evaluated)
|
||||||
|
elif self._operator == self.OPERATOR_OR:
|
||||||
|
return any(evaluated)
|
||||||
|
elif self._operator == self.OPERATOR_XOR:
|
||||||
|
result = evaluated[0]
|
||||||
|
for v in evaluated[1:]:
|
||||||
|
result ^= v
|
||||||
|
return result
|
||||||
|
elif self._operator == self.OPERATOR_INVERT:
|
||||||
|
return not evaluated[0]
|
||||||
|
return False
|
||||||
164
src/mab/filters/text.py
Normal file
164
src/mab/filters/text.py
Normal file
@@ -0,0 +1,164 @@
|
|||||||
|
from .base import BaseEventFilter
|
||||||
|
|
||||||
|
from nio import MatrixRoom, Event
|
||||||
|
|
||||||
|
class TextFilter(BaseEventFilter):
|
||||||
|
"""
|
||||||
|
This filter returns True if the event contains `body` attribute.
|
||||||
|
`body` attribute contains unformatted text, string.
|
||||||
|
"""
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
def __repr__(self):
|
||||||
|
return "TextFilter()"
|
||||||
|
|
||||||
|
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||||
|
return hasattr(event, "body") and type(event.body) is str # type: ignore
|
||||||
|
|
||||||
|
class FormattedTextFilter(BaseEventFilter):
|
||||||
|
"""
|
||||||
|
This filter returns True if the event contains valid `formatted_body`
|
||||||
|
attribute. `formatted_body` attribute contains formatted text, string.
|
||||||
|
"""
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
def __repr__(self):
|
||||||
|
return "FormattedTextFilter()"
|
||||||
|
|
||||||
|
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||||
|
return hasattr(event, "formatted_body") and type(event.formatted_body) is str # type: ignore
|
||||||
|
|
||||||
|
class TextContainsFilter(BaseEventFilter):
|
||||||
|
"""
|
||||||
|
This filter returns True if the `event.body` contains `needle`
|
||||||
|
substring (or any of neddle from the list). The check will be case
|
||||||
|
insensetive if `any_case` is True.
|
||||||
|
"""
|
||||||
|
def __init__(self, needle: str | list[str], *, any_case: bool = True):
|
||||||
|
super().__init__()
|
||||||
|
if type(needle) is str:
|
||||||
|
needle = [needle]
|
||||||
|
self._any_case = any_case
|
||||||
|
if self._any_case:
|
||||||
|
self._needle = [s.lower() for s in needle]
|
||||||
|
else:
|
||||||
|
self._needle = list(needle)
|
||||||
|
|
||||||
|
def __repr__(self):
|
||||||
|
return f"TextContainsFilter({repr(self._needle)}, any_case={repr(self._any_case)})"
|
||||||
|
|
||||||
|
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||||
|
try:
|
||||||
|
body = event.body.lower() if self._any_case else event.body # type: ignore
|
||||||
|
for n in self._needle:
|
||||||
|
if n in body:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
except:
|
||||||
|
return False
|
||||||
|
|
||||||
|
class TextStartsWithFilter(BaseEventFilter):
|
||||||
|
"""
|
||||||
|
This filter returns True if the `event.body` starts with `substring` (or
|
||||||
|
any of substrings from the list). The check will be case insensetive if
|
||||||
|
`any_case` is True.
|
||||||
|
"""
|
||||||
|
def __init__(self, substring: str | list[str], *, any_case: bool = True):
|
||||||
|
super().__init__()
|
||||||
|
if type(substring) is str:
|
||||||
|
substring = [substring]
|
||||||
|
self._any_case = any_case
|
||||||
|
if self._any_case:
|
||||||
|
self._substring = [s.lower() for s in substring]
|
||||||
|
else:
|
||||||
|
self._substring = list(substring)
|
||||||
|
|
||||||
|
def __repr__(self):
|
||||||
|
return f"TextStartsWithFilter({repr(self._substring)}, any_case={repr(self._any_case)})"
|
||||||
|
|
||||||
|
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||||
|
try:
|
||||||
|
body = event.body.lower() if self._any_case else event.body # type: ignore
|
||||||
|
for s in self._substring:
|
||||||
|
if body.startswith(s):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
except:
|
||||||
|
return False
|
||||||
|
|
||||||
|
class TextEndsWithFilter(BaseEventFilter):
|
||||||
|
"""
|
||||||
|
This filter returns True if the `event.body` ends with `substring` (or
|
||||||
|
any of substrings from the list). The check will be case insensetive if
|
||||||
|
`any_case` is True.
|
||||||
|
"""
|
||||||
|
def __init__(self, substring: str | list[str], *, any_case: bool = True):
|
||||||
|
super().__init__()
|
||||||
|
if type(substring) is str:
|
||||||
|
substring = [substring]
|
||||||
|
self._any_case = any_case
|
||||||
|
if self._any_case:
|
||||||
|
self._substring = [s.lower() for s in substring]
|
||||||
|
else:
|
||||||
|
self._substring = list(substring)
|
||||||
|
|
||||||
|
def __repr__(self):
|
||||||
|
return f"TextEndsWithFilter({repr(self._substring)}, any_case={repr(self._any_case)})"
|
||||||
|
|
||||||
|
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||||
|
try:
|
||||||
|
body = event.body.lower() if self._any_case else event.body # type: ignore
|
||||||
|
for s in self._substring:
|
||||||
|
if body.endswith(s):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
except:
|
||||||
|
return False
|
||||||
|
|
||||||
|
class TextCommandFilter(BaseEventFilter):
|
||||||
|
"""
|
||||||
|
This filter returns True if all conditions are True:
|
||||||
|
1. `event.body` contains at least `min_args + 1` words after split()
|
||||||
|
2. `event.body` conrains at most `max_args + 1` words after split()
|
||||||
|
3. First element of splitted `event.body` starts with `prefix`
|
||||||
|
4. First element of splitted `event.body` (after stripping `prefix`)
|
||||||
|
starts with any of strings in `verbs` list (case-insensitive)
|
||||||
|
|
||||||
|
Remarks:
|
||||||
|
- If this filter is satified, then it will set a new attribute for the
|
||||||
|
event: `event.command_args: list[str]`. You may use this attribute in
|
||||||
|
your callback for this event.
|
||||||
|
"""
|
||||||
|
def __init__(self, verbs: str | list[str], min_args: int = 0, max_args: int | None = None, prefix: str = "!"):
|
||||||
|
super().__init__()
|
||||||
|
if type(verbs) is str:
|
||||||
|
verbs = [verbs]
|
||||||
|
verbs = [v.lower() for v in verbs]
|
||||||
|
self._verbs = verbs
|
||||||
|
self._min_args = min_args
|
||||||
|
self._max_args = max_args
|
||||||
|
self._prefix = prefix
|
||||||
|
|
||||||
|
def __repr__(self):
|
||||||
|
return f"TextCommandFilter({repr(self._verbs)}, {repr(self._min_args)}, {repr(self._max_args)}, {repr(self._prefix)})"
|
||||||
|
|
||||||
|
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||||
|
try:
|
||||||
|
parts = [p.strip() for p in event.body.split() if p.strip()] # type: ignore
|
||||||
|
args_count = len(parts) - 1
|
||||||
|
if args_count < self._min_args:
|
||||||
|
return False
|
||||||
|
if self._max_args is not None and args_count > self._max_args:
|
||||||
|
return False
|
||||||
|
if not parts[0].startswith(self._prefix):
|
||||||
|
return False
|
||||||
|
cmd = parts[0][len(self._prefix):].lower()
|
||||||
|
for verb in self._verbs:
|
||||||
|
if cmd == verb:
|
||||||
|
setattr(event, "command_args", parts[1:])
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
except:
|
||||||
|
return False
|
||||||
@@ -20,4 +20,26 @@ class MatrixBotConfig:
|
|||||||
"""Enable automatic verification of all known devices on related errors"""
|
"""Enable automatic verification of all known devices on related errors"""
|
||||||
|
|
||||||
auto_join_any_room_on_invite: bool = True
|
auto_join_any_room_on_invite: bool = True
|
||||||
"""Enable automatic join for invites (implemented by automatically adding internal callback on bot creation)"""
|
"""Enable automatic join for invites (implemented by automatically adding internal callback on bot creation)"""
|
||||||
|
|
||||||
|
auto_ffprobe: bool = True
|
||||||
|
"""Enable automatic `ffprobe` execution to get some media properties"""
|
||||||
|
|
||||||
|
auto_ffmpeg: bool = True
|
||||||
|
"""Enabled automatic `ffmpeg` execution to generate some media (e.g. thumbnails)"""
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class VideoFileProperties:
|
||||||
|
"""Properties of the video for sending"""
|
||||||
|
|
||||||
|
width_px: int
|
||||||
|
"""Width in pixels"""
|
||||||
|
|
||||||
|
height_px: int
|
||||||
|
"""Height in pixels"""
|
||||||
|
|
||||||
|
duration_ms: int
|
||||||
|
"""Duration in milliseconds"""
|
||||||
|
|
||||||
|
thumbnail: Path | str | bytes | None = None
|
||||||
|
"""Path to the thumbnail or the raw JPEG thumbnail data"""
|
||||||
98
src/mab/utils.py
Normal file
98
src/mab/utils.py
Normal file
@@ -0,0 +1,98 @@
|
|||||||
|
"""Utility things"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from .types import *
|
||||||
|
|
||||||
|
class Utils:
|
||||||
|
@staticmethod
|
||||||
|
async def get_video_props_using_ffprobe(path: Path | str) -> VideoFileProperties:
|
||||||
|
"""Returns video properties for file at `path`.
|
||||||
|
|
||||||
|
Consider using `asyncio.timeout` to prevent deadlocking.
|
||||||
|
"""
|
||||||
|
# prepare arguments for ffprobe
|
||||||
|
args = [
|
||||||
|
"ffprobe",
|
||||||
|
"-v", "error",
|
||||||
|
"-select_streams", "v:0",
|
||||||
|
"-show_entries", "stream=width,height:format=duration",
|
||||||
|
"-of", "json",
|
||||||
|
str(path)
|
||||||
|
]
|
||||||
|
# start the process
|
||||||
|
proc = await asyncio.create_subprocess_exec(
|
||||||
|
*args,
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE,
|
||||||
|
)
|
||||||
|
# communicate `stdout` and `stderr`
|
||||||
|
try:
|
||||||
|
stdout, stderr = await proc.communicate()
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
if proc.returncode is None:
|
||||||
|
proc.kill()
|
||||||
|
await proc.wait()
|
||||||
|
raise
|
||||||
|
# error
|
||||||
|
if proc.returncode != 0:
|
||||||
|
raise RuntimeError(f"`ffprobe` exits with code {proc.returncode}: {stderr.decode(errors="replace")}")
|
||||||
|
# success
|
||||||
|
result = json.loads(stdout)
|
||||||
|
return VideoFileProperties(
|
||||||
|
width_px=int(result["streams"][0]["width"]),
|
||||||
|
height_px=int(result["streams"][0]["height"]),
|
||||||
|
duration_ms=int(float(result["format"]["duration"]) * 1000)
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
async def get_video_thumbnail(path: Path | str, *, timestamp: float = 0, max_size: int = 300) -> bytes:
|
||||||
|
"""Generates video thumbnail for file at `path`.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- path - path to the video
|
||||||
|
- timestamp - timestamp to use as the thumbnail
|
||||||
|
- max_size - max width/height
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Raw bytes of JPEG thumbnail
|
||||||
|
"""
|
||||||
|
# prepare arguments for ffmpeg
|
||||||
|
args = [
|
||||||
|
"ffmpeg",
|
||||||
|
"-v", "error",
|
||||||
|
"-ss", "%.2f" % timestamp,
|
||||||
|
"-i", str(path),
|
||||||
|
"-map", "0:v:0",
|
||||||
|
"-vf",
|
||||||
|
(
|
||||||
|
f"scale='min(iw,{max_size})':'min(ih,{max_size})':"
|
||||||
|
"force_original_aspect_ratio=decrease"
|
||||||
|
),
|
||||||
|
"-frames:v", "1",
|
||||||
|
"-q:v", "3",
|
||||||
|
"-f", "image2pipe",
|
||||||
|
"-vcodec", "mjpeg",
|
||||||
|
"pipe:1"
|
||||||
|
]
|
||||||
|
# start the process
|
||||||
|
proc = await asyncio.create_subprocess_exec(
|
||||||
|
*args,
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE,
|
||||||
|
)
|
||||||
|
# communicate `stdout` and `stderr`
|
||||||
|
try:
|
||||||
|
stdout, stderr = await proc.communicate()
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
if proc.returncode is None:
|
||||||
|
proc.kill()
|
||||||
|
await proc.wait()
|
||||||
|
raise
|
||||||
|
# error
|
||||||
|
if proc.returncode != 0 or not stdout:
|
||||||
|
raise RuntimeError(f"`ffmpeg` exits with code {proc.returncode}: {stderr.decode(errors="replace")}", args)
|
||||||
|
# success
|
||||||
|
return stdout
|
||||||
Reference in New Issue
Block a user