From 4ded810ba97c0d721590148e0c55cd5e5de18951 Mon Sep 17 00:00:00 2001
From: tt2468 <tt2468@gmail.com>
Date: Thu, 30 Apr 2020 20:09:13 -0700
Subject: [PATCH] Add optional `fileFormat` and `fileCompressionRatio`
 parameters to `TakeSourceScreenshot`

---
 src/WSRequestHandler_Sources.cpp | 32 +++++++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/WSRequestHandler_Sources.cpp b/src/WSRequestHandler_Sources.cpp
index 579f76b4..29ee2882 100644
--- a/src/WSRequestHandler_Sources.cpp
+++ b/src/WSRequestHandler_Sources.cpp
@@ -1437,6 +1437,8 @@ RpcResponse WSRequestHandler::SetSourceFilterVisibility(const RpcRequest& reques
 * @param {String} `sourceName` Source name. Note that, since scenes are also sources, you can also provide a scene name.
 * @param {String (optional)} `embedPictureFormat` Format of the Data URI encoded picture. Can be "png", "jpg", "jpeg" or "bmp" (or any other value supported by Qt's Image module)
 * @param {String (optional)} `saveToFilePath` Full file path (file extension included) where the captured image is to be saved. Can be in a format different from `pictureFormat`. Can be a relative path.
+* @param {String (optional)} `fileFormat` Format to save the image file as. If not specified, tries to guess based on file extension.
+* @param {int (optional)} `fileCompressionQuality` Compression ratio between -1 and 100 to save the file with. -1 is automatic, 1 is smallest file/most compression, 100 is largest file/least compression. Varies with image type.
 * @param {int (optional)} `width` Screenshot width. Defaults to the source's base width.
 * @param {int (optional)} `height` Screenshot height. Defaults to the source's base height.
 *
@@ -1544,7 +1546,7 @@ RpcResponse WSRequestHandler::TakeSourceScreenshot(const RpcRequest& request) {
 
 		QByteArrayList supportedFormats = QImageWriter::supportedImageFormats();
 		if (!supportedFormats.contains(pictureFormat)) {
-			QString errorMessage = QString("Unsupported picture format: %1").arg(pictureFormat);
+			QString errorMessage = QString("unsupported picture format: %1").arg(pictureFormat);
 			return request.failed(errorMessage.toUtf8());
 		}
 
@@ -1552,7 +1554,7 @@ RpcResponse WSRequestHandler::TakeSourceScreenshot(const RpcRequest& request) {
 		QBuffer buffer(&encodedImgBytes);
 		buffer.open(QBuffer::WriteOnly);
 		if (!sourceImage.save(&buffer, pictureFormat)) {
-			return request.failed("Embed image encoding failed");
+			return request.failed("embed image encoding failed");
 		}
 		buffer.close();
 
@@ -1569,7 +1571,31 @@ RpcResponse WSRequestHandler::TakeSourceScreenshot(const RpcRequest& request) {
 		QFileInfo filePathInfo(filePathStr);
 		QString absoluteFilePath = filePathInfo.absoluteFilePath();
 
-		if (!sourceImage.save(absoluteFilePath)) {
+		const char* fileFormat;
+		if (request.hasField("fileFormat")) {
+			fileFormat = obs_data_get_string(request.parameters(), "fileFormat");
+			QByteArrayList supportedFormats = QImageWriter::supportedImageFormats();
+
+			if (!supportedFormats.contains(fileFormat)) {
+				QString errorMessage = QString("unsupported file format: %1").arg(fileFormat);
+				return request.failed(errorMessage.toUtf8());
+			}
+		}
+		else {
+			fileFormat = nullptr;
+		}
+
+		int fileCompressionQuality {-1};
+		if (request.hasField("fileCompressionQuality")) {
+			fileCompressionQuality = obs_data_get_int(request.parameters(), "fileCompressionQuality");
+            
+			if (fileCompressionQuality < -1 || fileCompressionQuality > 100) {
+				QString errorMessage = QString("compression quality out of range: %1").arg(fileCompressionQuality);
+				return request.failed(errorMessage.toUtf8());
+			}
+		}
+
+		if (!sourceImage.save(absoluteFilePath, fileFormat, fileCompressionQuality)) {
 			return request.failed("Image save failed");
 		}
 		obs_data_set_string(response, "imageFile", absoluteFilePath.toUtf8());