From 9071decce12d810c3ffefbd27f75914ecffd2fd2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9D=D0=B0=D1=81=D1=8B=D1=80=D0=BE=D0=B2=20=D0=AE=D1=80?=
 =?UTF-8?q?=D0=B8=D0=B9=20=D0=A0=D0=B8=D0=BD=D0=B0=D1=82=D0=BE=D0=B2=D0=B8?=
 =?UTF-8?q?=D1=87?= <yu.nasyrov@2gis.ru>
Date: Thu, 31 Oct 2024 16:11:06 +0500
Subject: [PATCH] add color_transfer in video and count streams

---
 probedata.go | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/probedata.go b/probedata.go
index 9c9fed2..74131ce 100644
--- a/probedata.go
+++ b/probedata.go
@@ -80,6 +80,7 @@ type Stream struct {
 	Level              int               `json:"level,omitempty"`
 	ColorRange         string            `json:"color_range,omitempty"`
 	ColorSpace         string            `json:"color_space,omitempty"`
+	ColorTransfer      string            `json:"color_transfer,omitempty"`
 	SampleFmt          string            `json:"sample_fmt,omitempty"`
 	SampleRate         string            `json:"sample_rate,omitempty"`
 	Channels           int               `json:"channels,omitempty"`
@@ -181,6 +182,42 @@ func (p *ProbeData) FirstAttachmentStream() *Stream {
 	return p.firstStream(StreamAttachment)
 }
 
+// CountVideoStreams returns the count of the video streams
+func (p *ProbeData) CountVideoStreams() uint32 {
+	return p.countStreams(StreamVideo)
+}
+
+// CountAudioStreams returns the count of the audio streams
+func (p *ProbeData) CountAudioStreams() uint32 {
+	return p.countStreams(StreamAudio)
+}
+
+// CountSubtitleStreams returns the count of the subtitle streams
+func (p *ProbeData) CountSubtitleStreams() uint32 {
+	return p.countStreams(StreamSubtitle)
+}
+
+// CountDataStreams returns the count of the data streams
+func (p *ProbeData) CountDataStreams() uint32 {
+	return p.countStreams(StreamData)
+}
+
+// CountAttachmentStreams returns the count of the attacment streams
+func (p *ProbeData) CountAttachmentStreams() uint32 {
+	return p.countStreams(StreamAttachment)
+}
+
+func (p *ProbeData) countStreams(streamType StreamType) uint32 {
+	count := uint32(0)
+	for _, v := range p.Streams {
+		if v.CodecType == string(streamType) {
+			count++
+		}
+	}
+
+	return count
+}
+
 func (p *ProbeData) firstStream(streamType StreamType) *Stream {
 	for _, s := range p.Streams {
 		if s == nil {