"""collection of thread classes for handling FFmpeg streams

"""

from __future__ import annotations

from copy import deepcopy

import re, os

from threading import Thread, Condition, Lock, Event

from io import TextIOBase, TextIOWrapper

from time import sleep, time

from tempfile import TemporaryDirectory

from queue import Empty, Full, Queue

from math import ceil

import logging

logger = logging.getLogger("ffmpegio")

from .utils.avi import AviReader

from .utils.log import extract_output_stream as _extract_output_stream

from .errors import FFmpegError

# fmt:off

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

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

# fmt:on

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()

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 self

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, start=None, block=True, timeout=None):

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:

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 timeout - time()

# 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:

# 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 as e:

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, nmin=None, queuesize=None):

super().__init__()

self.stdout = stdout #:readable stream: data source

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

self.itemsize = None #:int: number of bytes per time sample

self._queue = Queue(queuesize or 0) # inter-thread data I/O

self._carryover = None # extra data that was not previously read by user

self._collect = True

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._collect = False

try:

self._queue.get_nowait()

except:

pass

def join(self, timeout=None):

if self._queue.full():

if timeout:

self._queue.not_full.wait(timeout)

if self._queue.full():

return

else:

with self._queue.mutex:

self._queue.queue.clear()

# if queue is full,

super().join(timeout)

def __enter__(self):

self.start()

return self

def __exit__(self, *_):

self.stdout.close()

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

return self

def run(self):

blocksize = (

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

) * self.itemsize

while True:

try:

data = self.stdout.read(blocksize)

except:

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

break

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

if not data:

if self.stdout.closed: # just in case

break

else:

break

if self._collect: # True until self.cooloff

self._queue.put(data)

# print(f"reader thread: queued samples")

def read(self, n=-1, timeout=None):

"""read n samples

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

:type n: int, optional

:param timeout: timeout in seconds, defaults to None

:type timeout: float, optional

:return: n*itemsize bytes

:rtype: bytes

"""

# wait till matching line is read by the thread

block = self.is_alive() and n != 0

if timeout is not None:

timeout = time() + timeout

arrays = []

n_new = max(n, -n)

# grab any leftover data from previous read

if self._carryover:

arrays = [self._carryover]

if n_new != 0:

n_new -= len(self._carryover) // self.itemsize

self._carryover = None

# loop till enough data are collected

nreads = 1 if n <= 0 else max(n_new, 0)

nr = 0

while True:

tout = timeout and timeout - time()

if tout <= 0:

break

try:

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

self._queue.task_done()

arrays.append(b)

except Empty:

break

nr += len(b) // self.itemsize

if nr >= nreads: # enough read

if n < 0:

block = False # keep reading until queue is empty

else:

break

# combine all the data and return requested amount

if not len(arrays):

return b""

all_data = b"".join(arrays)

if n <= 0:

return all_data

nbytes = self.itemsize * n

if len(all_data) > nbytes:

self._carryover = all_data[nbytes:]

return all_data[:nbytes]

def read_all(self, timeout=None):

# wait till matching line is read by the thread

if timeout is not None:

timeout = time() + timeout

arrays = arrays = [self._carryover] if self._carryover else []

self._carryover = None

# loop till enough data are collected

while not self.is_alive() or timeout and timeout > time():

try:

data = self._queue.get(self.is_alive(), timeout and timeout - time())

self._queue.task_done()

arrays.append(data)

except Empty:

break

# combine all the data and return requested amount

if not len(arrays):

return b""

return b"".join(arrays)

class WriterThread(Thread):

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

:param stdin: stream to write data to

:type stdin: writable stream

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

:type queuesize: int, optional

"""

def __init__(self, stdin, queuesize=None):

super().__init__()

self.stdin = stdin #:writable stream: data sink

self._queue = Queue(queuesize or 0) # inter-thread data I/O

def join(self, timeout=None):

# close the stream if not already closed

self.stdin.close()

# 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)

def __enter__(self):

self.start()

return self

def __exit__(self, *_):

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

return self

def run(self):

while True:

# get next data block

data = self._queue.get()

self._queue.task_done()

if data is None:

break

# print(f"writer thread: received {data.shape[0]} samples to write")

try:

nbytes = self.stdin.write(data)

# print(f"writer thread: written {nbytes} written")

except:

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

break

if not nbytes and self.stdin.closed: # just in case

break

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

if not self.is_alive():

raise ThreadNotActive("WriterThread is not running")

data = self._queue.put(data, timeout)

class AviReaderThread(Thread):

class InvalidAviStream(FFmpegError): ...

def __init__(self, queuesize=None):

super().__init__()

self.reader = AviReader() #:utils.avi.AviReader: AVI demuxer

self.streamsready = Event() #:Event: Set when received stream header info

self.rates = None # :dict(int:int|Fraction)

self._queue = Queue(queuesize or 0) # inter-thread data I/O

self._ids = None #:dict(int:int): stream indices

self._nread = None #:dict(int:int): number of samples read/stream

self._carryover = (

None #:dict(int:ndarray) extra data that was not previously read by user

)

@property

def streams(self):

return self.reader.streams if self.streamsready else None

def start(self, stdout, use_ya=None):

self._args = (stdout, use_ya)

super().start()

def join(self, timeout=None):

# if queue is full,

super().join(timeout)

def __bool__(self):

"""True if FFmpeg stdout stream is still open or there are more frames in the buffer"""

return self.is_alive() or not self._queue.empty()

# def __enter__(self):

# self.start()

# return self

# def __exit__(self, *_):

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

# return self

def run(self):

reader = self.reader

try:

# start the AVI reader to process stdout byte stream

reader.start(*self._args)

# initialize the stream properties

self._ids = ids = [i for i in reader.streams]

self._nread = {k: 0 for k in ids}

self.rates = {

k: v["frame_rate"] if v["type"] == "v" else v["sample_rate"]

for k, v in reader.streams.items()

}

except Exception as e:

logger.critical(e)

return

finally:

self.streamsready.set()

reader = self.reader

for id, data in reader:

self._queue.put((id, data))

self._queue.put(None) # end of stream

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

"""wait till stream is ready to be read

:param timeout: timeout in seconds, defaults to None (waits indefinitely)

:type timeout: float, optional

:raises InvalidAviStream: if thread has been terminated before stream header info was read

:return: tuple of stream specifier and data array

:rtype: (str, object)

"""

flag = self.streamsready.wait(timeout)

if not (flag or self.is_alive()):

raise self.InvalidAviStream(

"No stream header info was found in FFmpeg's AVI stream."

)

return flag

def readchunk(self, timeout=None):

"""read the next avi chunk

:param timeout: timeout in seconds, defaults to None (waits indefinitely)

:type timeout: float, optional

:raises TimeoutError: if terminated due to timeout

:return: tuple of stream specifier and data array

:rtype: (str, object)

"""

# wait till matching line is read by the thread

tend = timeout and time() + timeout

# if stream header not received in time, raise error

if not self.wait(timeout):

raise TimeoutError("timed out waiting for the stream headers")

block = self.is_alive()

# if any leftover data available, return the first one

if self._carryover is not None:

(id, data) = next(

((k, v) for k, v in self._carryover.items() if v is not None)

)

self._carryover[id] = None

if all((k for k, v in self._carryover.items() if v is None)):

self._carryover = None

return self.reader.streams[id]["spec"], self.reader.from_bytes(id, data)

# get next chunk

try:

if timeout is not None:

timeout = tend - time()

assert timeout > 0

chunk = self._queue.get(block, timeout)

if chunk is None:

raise ThreadNotActive("reached end-of-stream")

id, data = chunk

except Empty:

raise TimeoutError("timed out waiting for next chunk")

self._queue.task_done()

return self.reader.streams[id]["spec"], self.reader.from_bytes(id, data)

def find_id(self, ref_stream):

self.wait()

try:

return next(

(k for k, v in self.reader.streams.items() if v["spec"] == ref_stream)

)

except:

ValueError(f"{ref_stream} is not a valid stream specifier")

def read(self, n=-1, ref_stream=None, timeout=None):

"""read data from all streams

:param n: number of samples, negate to non-blocking, defaults to -1

:type n: int, optional

:param ref_stream: stream specifier to count the samples,

defaults to None (first stream)

:type ref_stream: str, optional

:param timeout: timeout in seconds, defaults to None (waits indefinitely)

:type timeout: float, optional

:raises TimeoutError: if terminated due to timeout

:return: tuple of stream specifier and data array

:return: dict of data object keyed by stream specifier string, each data object is

created by `bytes_to_video` or `bytes_to_image` plugin hook. If all frames

have been read, dict items would be all empty

:rtype: dict(spec:str, object)

"""

# wait till matching line is read by the thread

block = self.is_alive() and n != 0

tend = timeout and (time() + timeout)

# if stream header not received in time, raise error

if not self.wait(timeout):

raise TimeoutError("timed out waiting for the stream headers")

# get the reference stream id

if ref_stream is None:

ref_stream = self._ids[0]

else:

ref_stream = self.find_id(ref_stream)

# identify how many samples are needed for each stream

nref = max(n, -n)

tref = (self._nread[ref_stream] + nref) / self.rates[ref_stream]

n_need = {

k: ceil(tref * self.rates[k]) - self._nread[k] if k != ref_stream else nref

for k in self._ids

}

nremain = deepcopy(n_need)

# initialize output arrays

arrays = {k: [] for k in self._ids}

itemsizes = self.reader.itemsizes

# grab any leftover data from previous read

if self._carryover is not None:

for k, v in self._carryover.items():

if v is not None:

arrays[k] = [v]

nremain[k] -= len(v) // itemsizes[k]

self._carryover = None

# loop till enough data are collected

while any((v > 0 for k, v in nremain.items() if n_need[k] > 0)):

try:

if timeout:

timeout = tend - time()

if timeout <= 0:

break

chunk = self._queue.get(block, timeout)

if chunk is None:

break

k, data = chunk

self._queue.task_done()

arrays[k].append(data)

nremain[k] -= len(data) // itemsizes[k]

except Empty:

break

def combine(id, array, n, nr):

# combine all the data and return requested amount

if not len(array):

return (id, None, None)

all_data = b"".join(array)

nbytes = n * itemsizes[id]

return (

(id, all_data, None)

if nr >= 0

else (id, all_data[:nbytes], all_data[nbytes:])

)

ids, data, excess = zip(

*(

combine(id, array, n_need[id], nremain[id])

for id, array in arrays.items()

)

)

# any excess samples, store as a _carryover dict

if any((sdata is not None for sdata in excess)):

self._carryover = {id: sdata for id, sdata in zip(ids, excess)}

# final formatting of data

out = {}

for id, sdata in zip(ids, data):

info = self.reader.streams[id]

spec = info["spec"]

if sdata is None:

out[spec] = self.reader.from_bytes(id, b"")

else:

self._nread[id] += len(sdata) // itemsizes[id]

out[spec] = self.reader.from_bytes(id, sdata)

return out

def readall(self, timeout=None):

# wait till matching line is read by the thread

if timeout is not None:

timeout = time() + timeout

# if stream header not received in time, raise error

if not self.wait(timeout):

raise TimeoutError("timed out waiting for the stream headers")

# initialize output arrays

arrays = {k: [] for k in self._ids}

itemsizes = self.reader.itemsizes

# grab any leftover data from previous read

if self._carryover is not None:

for k, v in self._carryover.items():

if v is not None:

arrays[k] = [v]

self._nread[k] += len(v) // itemsizes[k]

self._carryover = None

# loop till enough data are collected

while True:

try:

chunk = self._queue.get(self.is_alive(), timeout and timeout - time())

if chunk is None:

break # end of stream

k, data = chunk

self._queue.task_done()

arrays[k].append(data)

self._nread[k] += len(data) // itemsizes[k]

except Empty:

break

# final formatting of data

out = {}

for id, sdata in arrays.items():

info = self.reader.streams[id]

spec = info["spec"]

out[spec] = self.reader.from_bytes(

id, b"" if sdata is None else b"".join(sdata)

)

return out