From 37ad1b7cb6b6aeb1de0b7a17216719311a9b5407 Mon Sep 17 00:00:00 2001
From: RogerSelwyn <rogers@selwyn.org.uk>
Date: Sat, 12 Feb 2022 17:48:13 +0000
Subject: [PATCH] Add ability to stop dowloading of attachments (to increase
 performance)

---
 custom_components/o365/const.py  |  3 +++
 custom_components/o365/sensor.py | 23 +++++++++++------------
 custom_components/o365/utils.py  |  9 ++++++---
 3 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/custom_components/o365/const.py b/custom_components/o365/const.py
index 757a96b..9d4a992 100644
--- a/custom_components/o365/const.py
+++ b/custom_components/o365/const.py
@@ -55,6 +55,7 @@ class EventResponse(Enum):
 CONF_CLIENT_ID = "client_id"
 CONF_CLIENT_SECRET = "client_secret"  # nosec
 CONF_DEVICE_ID = "device_id"
+CONF_DOWNLOAD_ATTACHMENTS = "download_attachments"
 CONF_EMAIL_SENSORS = "email_sensor"
 CONF_IGNORE_AVAILABILITY = "ignore_availability"
 CONF_OFFSET = "offset"
@@ -125,6 +126,7 @@ class EventResponse(Enum):
         vol.Optional(CONF_MAIL_FOLDER): cv.string,
         vol.Optional(CONF_MAX_ITEMS, default=5): int,
         vol.Optional(CONF_IS_UNREAD): bool,
+        vol.Optional(CONF_DOWNLOAD_ATTACHMENTS): bool,
     }
 )
 STATUS_SENSOR = vol.Schema(
@@ -143,6 +145,7 @@ class EventResponse(Enum):
         vol.Optional(CONF_IS_UNREAD): bool,
         vol.Exclusive(CONF_SUBJECT_CONTAINS, "subject_*"): cv.string,
         vol.Exclusive(CONF_SUBJECT_IS, "subject_*"): cv.string,
+        vol.Optional(CONF_DOWNLOAD_ATTACHMENTS): bool,
     }
 )
 CONFIG_SCHEMA = vol.Schema(
diff --git a/custom_components/o365/sensor.py b/custom_components/o365/sensor.py
index a809f66..3c6bc38 100644
--- a/custom_components/o365/sensor.py
+++ b/custom_components/o365/sensor.py
@@ -5,11 +5,11 @@
 
 from homeassistant.helpers.entity import Entity
 
-from .const import (CONF_EMAIL_SENSORS, CONF_HAS_ATTACHMENT, CONF_IMPORTANCE,
-                    CONF_IS_UNREAD, CONF_MAIL_FOLDER, CONF_MAIL_FROM,
-                    CONF_MAX_ITEMS, CONF_NAME, CONF_QUERY_SENSORS,
-                    CONF_STATUS_SENSORS, CONF_SUBJECT_CONTAINS,
-                    CONF_SUBJECT_IS, DOMAIN)
+from .const import (CONF_DOWNLOAD_ATTACHMENTS, CONF_EMAIL_SENSORS,
+                    CONF_HAS_ATTACHMENT, CONF_IMPORTANCE, CONF_IS_UNREAD,
+                    CONF_MAIL_FOLDER, CONF_MAIL_FROM, CONF_MAX_ITEMS,
+                    CONF_NAME, CONF_QUERY_SENSORS, CONF_STATUS_SENSORS,
+                    CONF_SUBJECT_CONTAINS, CONF_SUBJECT_IS, DOMAIN)
 from .utils import get_email_attributes
 
 _LOGGER = logging.getLogger(__name__)
@@ -89,6 +89,7 @@ def __init__(self, conf, mail_folder):
         """Initialise the O365 Sensor."""
         self.mail_folder = mail_folder
         self._name = conf.get(CONF_NAME)
+        self._download_attachments = conf.get(CONF_DOWNLOAD_ATTACHMENTS, True)
         self.max_items = conf.get(CONF_MAX_ITEMS, 5)
         self._state = 0
         self._attributes = {}
@@ -111,15 +112,13 @@ def extra_state_attributes(self):
 
     def update(self):
         """Update code."""
-        mails = list(
-            self.mail_folder.get_messages(
-                limit=self.max_items, query=self.query, download_attachments=True
-            )
+        data = self.mail_folder.get_messages(
+            limit=self.max_items, query=self.query, download_attachments=self._download_attachments
         )
-        attrs = [get_email_attributes(x) for x in mails]
+
+        attrs = [get_email_attributes(x, self._download_attachments) for x in data]
         attrs.sort(key=itemgetter("received"), reverse=True)
-        self._state = len(mails)
-        # self._attributes = {"data": attrs, "data_str_repr": json.dumps(attrs)}
+        self._state = len(attrs)
         self._attributes = {"data": attrs}
 
 
diff --git a/custom_components/o365/utils.py b/custom_components/o365/utils.py
index e0f5049..76e7982 100644
--- a/custom_components/o365/utils.py
+++ b/custom_components/o365/utils.py
@@ -72,9 +72,9 @@ def zip_files(filespaths, zip_name="archive.zip"):
     return zip_name
 
 
-def get_email_attributes(mail):
+def get_email_attributes(mail, download_attachments = True):
     """Get the email attributes."""
-    return {
+    data = {
         "subject": mail.subject,
         "body": clean_html(mail.body),
         "received": mail.received.strftime(DATETIME_FORMAT),
@@ -85,8 +85,11 @@ def get_email_attributes(mail):
         "has_attachments": mail.has_attachments,
         "importance": mail.importance.value,
         "is_read": mail.is_read,
-        "attachments": [x.name for x in mail.attachments],
     }
+    if download_attachments:
+        data["attachments"] = [x.name for x in mail.attachments]
+
+    return data
 
 
 def format_event_data(event, calendar_id):