From 35cdef13400e7d370eaafa161b46e4c5ee29c25f Mon Sep 17 00:00:00 2001
From: JonnyWong16 <9099342+JonnyWong16@users.noreply.github.com>
Date: Sun, 20 Sep 2020 12:27:42 -0700
Subject: [PATCH] Add hdr attribute to media export

---
 plexpy/exporter.py | 17 ++++++++++-------
 plexpy/helpers.py  | 14 +++++++++++++-
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/plexpy/exporter.py b/plexpy/exporter.py
index 9ddf0c21..7fcaf844 100644
--- a/plexpy/exporter.py
+++ b/plexpy/exporter.py
@@ -110,6 +110,7 @@ MOVIE_ATTRS = {
         'videoProfile': None,
         'videoResolution': None,
         'width': None,
+        'hdr': lambda i: get_any_hdr(i, MOVIE_ATTRS['media']),
         'parts': {
             'accessible': None,
             'audioProfile': None,
@@ -389,6 +390,7 @@ EPISODE_ATTRS = {
         'videoProfile': None,
         'videoResolution': None,
         'width': None,
+        'hdr': lambda i: get_any_hdr(i, EPISODE_ATTRS['media']),
         'parts': {
             'accessible': None,
             'audioProfile': None,
@@ -876,7 +878,7 @@ MOVIE_LEVELS = {
         'locations', 'media.aspectRatio', 'media.audioChannels', 'media.audioCodec', 'media.audioProfile',
         'media.bitrate', 'media.container', 'media.duration', 'media.height', 'media.width',
         'media.videoCodec', 'media.videoFrameRate', 'media.videoProfile', 'media.videoResolution',
-        'media.optimizedVersion'
+        'media.optimizedVersion', 'media.hdr'
     ],
     5: [
         'media.parts.accessible', 'media.parts.exists', 'media.parts.file', 'media.parts.duration',
@@ -944,7 +946,12 @@ MEDIA_TYPES = {
 }
 
 
-def export(section_id=None, rating_key=None, file_format='json', level=1):
+def get_any_hdr(obj, root):
+    attrs = helpers.get_dict_value_by_path(root, 'parts.videoStreams.hdr')
+    media = helpers.get_attrs_to_dict(obj, attrs)
+    return any(vs.get('hdr') for p in media.get('parts', []) for vs in p.get('videoStreams', []))
+
+
     timestamp = helpers.timestamp()
 
     level = helpers.cast_to_int(level)
@@ -1022,16 +1029,12 @@ def export(section_id=None, rating_key=None, file_format='json', level=1):
 
         export_attrs_list = []
         for attr in export_attrs_set:
-            split_attr = attr.split('.')
             try:
-                value = helpers.get_by_path(media_attrs, split_attr)
+                value = helpers.get_dict_value_by_path(media_attrs, attr)
             except KeyError:
                 logger.warn("Tautulli Exporter :: Unknown export attribute '%s', skipping...", attr)
                 continue
 
-            for _attr in reversed(split_attr):
-                value = {_attr: value}
-
             export_attrs_list.append(value)
 
         export_attrs = reduce(helpers.dict_merge, export_attrs_list)
diff --git a/plexpy/helpers.py b/plexpy/helpers.py
index 4ca630a4..bde7a4b6 100644
--- a/plexpy/helpers.py
+++ b/plexpy/helpers.py
@@ -1,4 +1,4 @@
-# -*- coding: utf-8 -*-
+# -*- coding: utf-8 -*-
 
 #  This file is part of Tautulli.
 #
@@ -1267,14 +1267,26 @@ def flatten_tree(obj, key=''):
 # https://stackoverflow.com/a/14692747
 def get_by_path(root, items):
     """Access a nested object in root by item sequence."""
+    if isinstance(items, str):
+        items = items.split('.')
     return reduce(operator.getitem, items, root)
 
 
 def set_by_path(root, items, value):
     """Set a value in a nested object in root by item sequence."""
+    if isinstance(items, str):
+        items = items.split('.')
     get_by_path(root, items[:-1])[items[-1]] = value
 
 
+def get_dict_value_by_path(root, attr):
+    split_attr = attr.split('.')
+    value = get_by_path(root, split_attr)
+    for _attr in reversed(split_attr):
+        value = {_attr: value}
+    return value
+
+
 # https://stackoverflow.com/a/7205107
 def dict_merge(a, b, path=None):
     if path is None: