From ae37560a9b6695b0bf5078eaa4ac390c22fb383c Mon Sep 17 00:00:00 2001
From: FoxxMD <FoxxMD@users.noreply.github.com>
Date: Fri, 23 Feb 2024 12:50:46 -0500
Subject: [PATCH] feat: Provide colorette object to messageFormat

Enables users to use available colors based on `colorize` context of the pino-pretty instance
---
 Readme.md                          |  6 ++++--
 index.d.ts                         |  3 ++-
 lib/colors.js                      |  3 +++
 lib/utils/index.js                 |  1 +
 lib/utils/prettify-message.js      |  2 +-
 lib/utils/prettify-message.test.js | 34 ++++++++++++++++++++++++++++++
 6 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/Readme.md b/Readme.md
index 41a504a4..9b1f73ca 100644
--- a/Readme.md
+++ b/Readme.md
@@ -356,9 +356,11 @@ This option can also be defined as a `function` with this prototype:
 
 ```js
 {
-  messageFormat: (log, messageKey, levelLabel) => {
+  messageFormat: (log, messageKey, levelLabel, colors) => {
     // do some log message customization
-    return customized_message;
+    //
+    // `colors` is a Colorette object with colors enabled based on `colorize` option
+    return `This is a ${color.red('colorized')}, custom message: ${log[messageKey]}`;
   }
 }
 ```
diff --git a/index.d.ts b/index.d.ts
index 614d3750..fe1dd242 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -10,6 +10,7 @@ import { Transform } from 'stream';
 import { OnUnknown } from 'pino-abstract-transport';
 // @ts-ignore fall back to any if pino is not available, i.e. when running pino tests
 import { DestinationStream, Level } from 'pino';
+import * as Colorette from "colorette";
 
 type LogDescriptor = Record<string, unknown>;
 
@@ -205,7 +206,7 @@ declare function build(options: PrettyOptions_): PinoPretty.PrettyStream;
 
 declare namespace PinoPretty {
   type Prettifier = (inputData: string | object) => string;
-  type MessageFormatFunc = (log: LogDescriptor, messageKey: string, levelLabel: string) => string;
+  type MessageFormatFunc = (log: LogDescriptor, messageKey: string, levelLabel: string, colors: Colorette.Colorette) => string;
   type PrettyOptions = PrettyOptions_;
   type PrettyStream = Transform & OnUnknown;
   type ColorizerFactory = typeof colorizerFactory;
diff --git a/lib/colors.js b/lib/colors.js
index 5bb38042..08dc2a43 100644
--- a/lib/colors.js
+++ b/lib/colors.js
@@ -67,6 +67,7 @@ function plainColorizer (useOnlyCustomProps) {
   }
   customColoredColorizer.message = plain.message
   customColoredColorizer.greyMessage = plain.greyMessage
+  customColoredColorizer.colors = createColors({ useColor: false })
   return customColoredColorizer
 }
 
@@ -77,6 +78,7 @@ function coloredColorizer (useOnlyCustomProps) {
   }
   customColoredColorizer.message = colored.message
   customColoredColorizer.greyMessage = colored.greyMessage
+  customColoredColorizer.colors = availableColors
   return customColoredColorizer
 }
 
@@ -105,6 +107,7 @@ function customColoredColorizerFactory (customColors, useOnlyCustomProps) {
  * recognized.
  * @property {function} message Accepts one string parameter that will be
  * colorized to a predefined color.
+ * @property {Colorette.Colorette} colors Available color functions based on `useColor` (or `colorize`) context
  */
 
 /**
diff --git a/lib/utils/index.js b/lib/utils/index.js
index ec69d968..dedc6bc0 100644
--- a/lib/utils/index.js
+++ b/lib/utils/index.js
@@ -90,6 +90,7 @@ module.exports = {
  * contains the log message.
  * @param {string} levelLabel The name of the key in the `log` object that
  * contains the log level name.
+ * @param {Colorette.Colorette} colors Available color functions based on `colorize` context
  * @returns {string}
  *
  * @example
diff --git a/lib/utils/prettify-message.js b/lib/utils/prettify-message.js
index 7758f1d8..5b867bbc 100644
--- a/lib/utils/prettify-message.js
+++ b/lib/utils/prettify-message.js
@@ -54,7 +54,7 @@ function prettifyMessage ({ log, context }) {
     return colorizer.message(message)
   }
   if (messageFormat && typeof messageFormat === 'function') {
-    const msg = messageFormat(log, messageKey, levelLabel)
+    const msg = messageFormat(log, messageKey, levelLabel, colorizer.colors)
     return colorizer.message(msg)
   }
   if (messageKey in log === false) return undefined
diff --git a/lib/utils/prettify-message.test.js b/lib/utils/prettify-message.test.js
index 8faf4b12..fc668dc9 100644
--- a/lib/utils/prettify-message.test.js
+++ b/lib/utils/prettify-message.test.js
@@ -185,3 +185,37 @@ tap.test('`messageFormat` supports function definition', async t => {
   })
   t.equal(str, '--> localhost/test')
 })
+
+tap.test('`messageFormat` supports function definition with colorizer object', async t => {
+  const colorizer = getColorizer(true)
+  const str = prettifyMessage({
+    log: { level: 30, request: { url: 'localhost/test' }, msg: 'incoming request' },
+    context: {
+      ...context,
+      colorizer,
+      messageFormat: (log, messageKey, levelLabel, colors) => {
+        let msg = log[messageKey]
+        if (msg === 'incoming request') msg = `--> ${colors.red(log.request.url)}`
+        return msg
+      }
+    }
+  })
+  t.equal(str, '\u001B[36m--> \u001B[31mlocalhost/test\u001B[36m\u001B[39m')
+})
+
+tap.test('`messageFormat` supports function definition with colorizer object when no color is supported', async t => {
+  const colorizer = getColorizer(false)
+  const str = prettifyMessage({
+    log: { level: 30, request: { url: 'localhost/test' }, msg: 'incoming request' },
+    context: {
+      ...context,
+      colorizer,
+      messageFormat: (log, messageKey, levelLabel, colors) => {
+        let msg = log[messageKey]
+        if (msg === 'incoming request') msg = `--> ${colors.red(log.request.url)}`
+        return msg
+      }
+    }
+  })
+  t.equal(str, '--> localhost/test')
+})