"""collection of thread classes for handling FFmpeg streams"""

from __future__ import annotations

import logging

import os

import re

from io import TextIOBase, TextIOWrapper

from queue import Empty, Full, Queue

from shutil import copyfileobj

from tempfile import TemporaryDirectory

from threading import Condition, Event, Lock, Thread

from time import sleep, time

from typing import BinaryIO

from namedpipe import NPopen

from .errors import FFmpegError

from .utils.log import extract_output_stream as _extract_output_stream

logger = logging.getLogger("ffmpegio")

# fmt:off

__all__ = ['FFmpegError', 'ThreadNotActive', 'ProgressMonitorThread',

'LoggerThread', 'ReaderThread', 'WriterThread', 'Empty', 'Full']

# fmt:on

class NotEmpty(Exception):

"Exception raised by WriterThread.flush(timeout) if timedout."

pass

class ThreadNotActive(RuntimeError):

pass

class ProgressMonitorThread(Thread):

"""FFmpeg progress monitor class

:param callback: [description]

:type callback: function

:param cancel_fun: [description], defaults to None

:type cancel_fun: [type], optional

:param url: [description], defaults to None

:type url: [type], optional

:param timeout: [description], defaults to 10e-3

:type timeout: [type], optional

"""

def __init__(self, callback, cancelfun=None, url=None, timeout=10e-3):

if callback is None:

self.url = self.cancelfun = self._thread = None

else:

tempdir = None if url else TemporaryDirectory()

self.url = url or os.path.join(tempdir.name, "progress.txt")

self.cancelfun = cancelfun

super().__init__(args=(callback, tempdir, timeout))

self._stop_monitor = Event()

def start(self):

if self.url:

super().start()

def join(self, timeout=None):

if self.url:

self._stop_monitor.set()

super().join(timeout)

def __enter__(self):

self.start()

return self

def __exit__(self, exc_type, exc_value, traceback):

self.join()

return False

def run(self):

callback, tempdir, timeout = self._args

url = self.url

pattern = re.compile(r"(.+)?=(.+)")

logger.debug(f'[progress_monitor] monitoring "{url}"')

while not (self._stop_monitor.is_set() or os.path.isfile(url)):

sleep(timeout)

logger.debug("[progress_monitor] file found")

if not self._stop_monitor.is_set():

with open(url, "rt") as f:

last_mtime = None

def update(sleep=True):

d = {}

mtime = os.fstat(f.fileno()).st_mtime

new_data = mtime != last_mtime

if new_data:

lines = f.readlines()

for line in lines:

m = pattern.match(line)

if not m:

continue

if m[1] != "progress":

val = m[2].lstrip()

try:

val = int(val)

except:

try:

val = float(val)

except:

pass

d[m[1]] = val

else:

done = m[2] == "end"

try:

if callback(d, done) and self.cancelfun:

logger.debug(

"[progress_monitor] operation canceled by user agent"

)

self.cancelfun()

self.cancelfun = None

except Exception as e:

logger.critical(

f"[progress_monitor] user callback error:\n\n{e}"

)

elif sleep:

sleep(timeout)

while not self._stop_monitor.is_set():

last_mtime = update()

# one final update just in case FFmpeg termianted during sleep

update(False)

if tempdir is not None:

try:

tempdir.cleanup()

except:

pass

logger.debug("[progress_monitor] terminated")

class LoggerThread(Thread):

def __init__(self, stderr, echo=False) -> None:

self.stderr = stderr

self.logs = []

self._newline_mutex = Lock()

self.newline = Condition(self._newline_mutex)

self.echo = echo

super().__init__()

def __enter__(self):

self.start()

return self

def __exit__(self, *_):

self.stderr.close()

self.join() # will wait until stderr is closed

return False

def run(self):

logger.debug("[logger] starting")

stderr = self.stderr

if not stderr or stderr.closed:

logger.debug("[logger] exiting (stderr pipe not open)")

return

if not isinstance(stderr, TextIOBase):

stderr = self.stderr = TextIOWrapper(stderr, "utf-8")

while True:

try:

log = stderr.readline()

except:

# stderr stream closed/FFmpeg terminated, end the thread as well

break

if not log and stderr.closed:

break

log = log[:-1] # remove the newline

if not log:

sleep(0.001)

continue

if self.echo:

print(log)

logger.debug(log)

with self.newline:

self.logs.append(log)

self.newline.notify_all()

with self.newline:

self.stderr = None

self.newline.notify_all()

logger.debug("[logger] exiting")

def index(

self,

prefix: str,

start: int = 0,

block: bool = True,

timeout: float | None = None,

) -> int | None:

"""Return an index of the first log line which starts with the specified prefix

:param prefix: look for log lines starting with this string

:param start: log line index to start searching, defaults to 0

:param block: True to block until the specified log line appears, default is True

:param timeout: blocking timeout, defaults to None (wait indefinitely)

:return: index of the matching line of the LoggerThread.logs or None if none found

"""

start = int(start or 0)

with self.newline:

logs = self.logs[start:] if start else self.logs

try:

# check existing lines

return (

next((i for i, log in enumerate(logs) if log.startswith(prefix)))

+ start

)

except StopIteration:

if not self.is_alive():

raise ThreadNotActive("LoggerThread is not running")

# no wait mode

if not block:

raise ValueError("Specified line not found")

# wait till matching line is read by the thread

if timeout is not None:

timeout = time() + timeout

start = len(self.logs)

while True:

tout = timeout and max(timeout - time(), 0)

# wait till the next log update

if (tout is not None and tout < 0) or not self.newline.wait(tout):

raise TimeoutError("Specified line not found")

# FFmpeg could have been terminated without match

if self.stderr is None:

raise ValueError("Specified line not found")

# check the new lines

try:

return (

next(

(

i

for i, log in enumerate(self.logs[start:])

if log.startswith(prefix)

)

)

+ start

)

except StopIteration:

# still no match, update the starting position

start = len(self.logs)

def output_stream(self, file_id=0, stream_id=0, block=True, timeout=None):

try:

i = self.index(f"Output #{file_id}", block=block, timeout=timeout)

self.index(f" Stream #{file_id}:{stream_id}", i, block, timeout)

except ThreadNotActive as e:

raise e

except TimeoutError:

raise TimeoutError("Specified output stream not found")

except Exception:

raise ValueError("Specified output stream not found")

with self._newline_mutex:

return _extract_output_stream(self.logs, hint=i)

def join_and_raise(self, timeout: float | None = None):

"""wait till thread terminates and raise exception based on the log

:param timeout: specifying a timeout for the operation in seconds if present, defaults to None

:type timeout: float | None, optional

:raises e: FFmpegError only if log is present

Note: This method throws the exception regardless of the thread's status if log is available.

"""

self.join(timeout)

e = self.Exception

if e is not None:

raise e

@property

def Exception(self) -> FFmpegError | None:

"""Exception gathered from the current log or None if there is no log"""

return FFmpegError(self.logs) if len(self.logs) else None

class ReaderThread(Thread):

def __init__(

self,

stdout_or_pipe: BinaryIO | NPopen,

nmin: int | None = None,

queuesize: int | None = None,

itemsize: int | None = None,

retry_delay: float | None = None,

timeout: float | None = None,

):

super().__init__()

is_pipe = isinstance(stdout_or_pipe, NPopen)

self.pipe = stdout_or_pipe if is_pipe else None # readable named pipe

self.stdout = None if is_pipe else stdout_or_pipe #:readable stream

self.nmin = nmin #:positive int: expected minimum number of read()'s n arg (not enforced)

self.itemsize = itemsize or 2**20 #:int: number of bytes per time sample

self._queue = Queue(16 if queuesize is None else queuesize)

self._carryover: bytes | None = (

None #:bytes: extra data that was not previously read by user

)

self._halt = Event()

self._cooling = Event()

self._running = Event()

self._retry_delay = 0.01 if retry_delay is None else retry_delay

self._timeout = float(timeout) if timeout else None

def start(self):

if self.itemsize is None:

raise ValueError(

"Thread object's must have its itemsize property set with the expected sample/frame size in bytes"

)

super().start()

def cool_down(self):

# stop enqueue read samples

self._cooling.set()

def join(self, timeout=None):

if timeout is None:

timeout = self._timeout

if self.pipe is None:

self.stdout.close()

else:

if self.stdout is None:

# FFmpeg never opened the pipe, open it to release the runner from waiting

with open(self.pipe.path, "w"):

...

self.pipe.close()

# set flag to terminate the thread loop

self._cooling.set()

self._halt.set()

# if queue is full, the thread loop is likely stuck.

is_full = self._queue.full()

if is_full and timeout:

# if timeout is set, wait to see if dequeued by another thread

self._queue.not_full.wait(timeout)

is_full = self._queue.full()

if is_full:

raise Full("Cannot join reader thread because the queue is full.")

# if queue is full,

super().join(timeout)

def is_running(self):

return self._running.is_set()

def wait_till_running(self, timeout: float | None = None) -> bool:

return self._running.wait(timeout or self._timeout)

def __enter__(self):

self.start()

return self

def __exit__(self, *_):

self.join() # will wait until stdout is closed

return False

def run(self):

is_npipe = self.stdout is None

blocksize = (

self.nmin if self.nmin is not None else 1 if self.itemsize > 1024 else 1024

) * self.itemsize

if self._halt.is_set():

return

logger.debug("waiting for pipe to open")

if is_npipe:

self.stdout = self.pipe.wait()

stream = self.stdout

queue = self._queue

logger.debug("starting to read")

self._running.set()

while not self._cooling.is_set():

try:

data = stream.read(blocksize)

# logger.debug("read %d bytes", len(data))

except:

# stdout stream closed/FFmpeg terminated, end the thread as well

data = None

# print(f"reader thread: read {len(data)} bytes")

if data:

logger.debug("ReaderThread putting data into the queue")

while not self._cooling.is_set():

try:

queue.put(data, timeout=0.01)

break

except Full:

if self._cooling.is_set():

break

logger.debug("ReaderThread data in the reader queue")

elif stream.closed: # just in case

logger.info("ReaderThread no data, stream is closed, exiting")

self._cooling.set()

self._halt.set()

break

else:

# pause a bit then try again

# logger.info("ReaderThread no data, reader thread pausing")

sleep(self._retry_delay)

logger.debug("stopping to read")

logger.info("ReaderThread sending the sentinel")

while not self._halt.is_set():

try:

queue.put(None, timeout=0.01)

break

except Full:

if self._halt.is_set():

break

# cooling loop (no queuing, flush all read)

logger.info("ReaderThread enters cool-down mode")

while not self._halt.is_set():

stream.read(blocksize)

logger.info("ReaderThread exiting")

self._running.clear()

def read(self, n: int = -1, timeout: float | None = None) -> bytes:

"""read n samples

:param n: number of samples/frames to read, if non-positive, read all

(until the pipe is broken), defaults to -1

:param timeout: timeout in seconds, defaults to wait indefinitely

:return: n*itemsize bytes

"""

# no sample requested, return empty bytes object

if n == 0:

return b""

read_all = n < 0

# wait till matching line is read by the thread

if timeout is None:

timeout = self._timeout

if timeout is not None:

timeout = time() + timeout

arrays = []

m = n * self.itemsize # bytes needed

mread = 0 # bytes read

# grab any leftover data from previous read

if self._carryover:

mread = len(self._carryover)

arrays = [self._carryover]

m -= mread

self._carryover = None

# loop till enough data are collected

while read_all or m > 0:

tout = timeout and max(timeout - time(), 0)

block = self.is_alive() and timeout is None

try:

b = self._queue.get(block, tout or 0.01)

except Empty:

if not block:

break

else:

if b is None:

# encountered sentinel

break

self._queue.task_done()

arrays.append(b)

mr = len(b)

m -= mr

mread += mr

assert mr and (

read_all or tout is None or tout > 0

) # no more read time left

# combine all the data and return requested amount

all_data = b"".join(arrays)

nread = mread // self.itemsize # number of frames read

if n >= 0:

nread = min(nread, n) # adjust to number of frames needed

mbytes = nread * self.itemsize # number of bytes needed

# update carryover buffer

self._carryover = all_data[mbytes:] if mbytes < mread else None

# return retrieved bytes array

return all_data[:mbytes]

def read_all(self, timeout: float | None = None) -> bytes:

return self.read(-1, timeout)

def read_nowait(self, n: int = -1) -> bytes:

"""read at most n samples

:param n: number of samples/frames to read, if non-positive, read all

in the queue, defaults to -1

:return: <= n*itemsize bytes

"""

# no sample requested, return empty bytes object

if n == 0:

return b""

read_all = n < 0

# wait till matching line is read by the thread

arrays = []

m = n * self.itemsize # bytes needed

mread = 0 # bytes read

# grab any leftover data from previous read

if self._carryover:

mread = len(self._carryover)

arrays = [self._carryover]

m -= mread

self._carryover = None

# loop till enough data are collected

while read_all or m > 0:

try:

b = self._queue.get_nowait()

self._queue.task_done()

if b is None:

# sentinel

break

mr = len(b)

assert mr > 0 # just in case

arrays.append(b)

m -= mr

mread += mr

except Empty:

break

# combine all the data and return requested amount

all_data = b"".join(arrays)

nread = mread // self.itemsize # number of frames read

if n >= 0:

nread = min(nread, n) # adjust to number of frames needed

mbytes = nread * self.itemsize # number of bytes needed

# update carryover buffer

self._carryover = all_data[mbytes:] if mbytes < mread else None

# return retrieved bytes array

return all_data[:mbytes]

def qsize(self) -> int:

"""Return the approximate size of the queue.

Note, qsize() > 0 doesn't guarantee that a subsequent write() will not block,

nor will qsize() < maxsize guarantee that put() will not block.

"""

return self._queue.qsize()

def empty(self) -> bool:

"""Return True if the queue is empty, False otherwise.

If empty() returns False it doesn't guarantee that a subsequent call to

read() will not block.

"""

return self._queue.empty()

def full(self) -> bool:

"""Return True if the queue is full, False otherwise.

If full() returns True it doesn't guarantee that a subsequent call to

read() will not block.

"""

return self._queue.full()

class WriterThread(Thread):

"""a thread to write byte data to a writable stream

:param stdin: stream to write data to

:param queuesize: depth of a queue for inter-thread data transfer, defaults to None

:param timeout: maximum number of bytes to write at once, defaults to None (1048576 bytes)

"""

def __init__(

self,

stdin_or_pipe: BinaryIO | NPopen,

queuesize: int | None = None,

timeout: float | None = None,

):

super().__init__()

is_pipe = isinstance(stdin_or_pipe, NPopen)

self.pipe = stdin_or_pipe if is_pipe else None

self.stdin = None if is_pipe else stdin_or_pipe #:writable stream: data sink

self._queue = Queue(16 if queuesize is None else queuesize)

self._empty_cond = Condition()

self._empty = True

self._no_more = False # true if sentinel has been written to the queue

self._timeout = float(timeout) if timeout else None

def join(self, timeout: float | None = None):

if self.stdin is None:

# pipe not yet connected, open it to release the runner

with open(self.pipe.path, "rb"):

...

# if empty, queue a dummy item to wake up the thread

if self._queue.empty():

self._queue.put(None)

# if queue is full,

super().join(timeout or self._timeout)

def closed(self) -> bool:

"""``True`` if thread no longer accepts data to write"""

return self._no_more

def __enter__(self):

self.start()

return self

def __exit__(self, *_):

self.join() # will wait until stdout is closed

return False

def run(self):

if self.stdin is None:

self.stdin = self.pipe.wait()

stream = self.stdin

queue = self._queue

while True:

# get next data block

logger.debug("WriterThread getting data to the queue")

try:

data = queue.get_nowait()

except Empty:

# if empty, set the flag and block

with self._empty_cond:

self._empty = True

self._empty_cond.notify_all()

data = queue.get()

logger.debug("WriterThread getting data from the queue")

queue.task_done()

if data is None:

logger.debug("WriterThread: received a sentinel to stop the writer")

break

else:

logger.debug("WriterThread: writing %d bytes", len(data))

try:

nwritten = 0

nwritten = stream.write(data)

logger.debug("WriterThread: written %d written", nwritten)

except Exception as e:

# stdout stream closed/FFmpeg terminated, end the thread as well

logger.debug("WriterThread exception: %s", e)

break

if not nwritten and stream.closed: # just in case

logger.debug("WriterThread: somethin' else happened")

break

# set flag to prevent any more writes

with self._empty_cond:

self._no_more = True

# close the pipe/stream

if self.pipe is not None:

self.pipe.close()

elif not self.stdin.closed:

self.stdin.close()

# completely flush the queue

# check if queue has any remaining items

not_empty = True

while True:

try:

queue.get_nowait()

except Empty:

not_empty = False

break

# if queue was not empty, notify its empty state now

if not_empty:

with self._empty_cond:

self._empty = True

self._empty_cond.notify_all()

logger.info("WriterThread exiting")

def write(self, data, timeout=None):

with self._empty_cond:

if self._no_more:

if data is None:

return

else:

raise ThreadNotActive("WriterThread is no longer running")

if data is None:

self._no_more = True

self._queue.put(data, timeout != 0, timeout or self._timeout)

self._empty = False

def flush(self, timeout: float | None = None):

"""block until the write buffer is emptied

:param timeout: a timeout for blocking in seconds, or fractions

thereof, defaults to None, to wait until empty

:raise NotEmpty: if a timeout is set, and the buffer is not emptied in time

"""

with self._empty_cond:

if not (

self._no_more

or self._empty

or self._empty_cond.wait(timeout or self._timeout)

):

raise NotEmpty()

def qsize(self) -> int:

"""Return the approximate size of the queue.

Note, qsize() > 0 doesn't guarantee that a subsequent write() will not block

"""

return self._queue.qsize()

def empty(self) -> bool:

"""Return True if the write queue is empty, False otherwise.

If empty() returns True it doesn't guarantee that a subsequent call to

write() will not block.

"""

return self._queue.empty()

def full(self) -> bool:

"""Return True if the queue is full, False otherwise.

If full() returns False it doesn't guarantee that a subsequent call to

write() will not block.

"""

return self._queue.full()

class CopyFileObjThread(Thread):

"""run shutil.copyfileobj in the thread

:param fsrc: source file object

:param fout: destination file object

:param length: The integer length, if given, is the buffer size. In particular, a negative length

value means to copy the data without looping over the source data in chunks;

defaults to 0; the data is read in chunks to avoid uncontrolled memory consumption.

:param auto_close: True for the thread to close fsrc and fdst after copy,

defaults to False

Thread terminates when the copy operation is completed.

Note that if the current file position of the fsrc object is not 0,

only the contents from the current file position to the end of the file will be copied.

"""

def __init__(

self,

fsrc: BinaryIO | NPopen,

fdst: BinaryIO | NPopen,

length: int = 0,

*,

auto_close: bool = False,

):

super().__init__()

self._fsrc = fsrc

self._fdst = fdst

self.length = length

self.auto_close = auto_close

def __enter__(self):

self.start()

return self

def __exit__(self, *_):

self.join()

return False

def run(self):

src_is_namedpipe = isinstance(self._fsrc, NPopen)

src = self._fsrc.wait() if src_is_namedpipe else self._fsrc

dst_is_namedpipe = isinstance(self._fdst, NPopen)

dst = self._fdst.wait() if dst_is_namedpipe else self._fdst

try:

copyfileobj(src, dst, self.length)

except:

# TODO - test the behavior when FFmpeg is prematurely terminated

logger.warning("CopyFileObjThread runner failed to complete the job.")

if self.auto_close:

src.close()

dst.close()