refactor: convert more functions to async

custom_components/mail_and_packages/helpers.py

@@ -33,7 +33,7 @@
     CONF_USERNAME,
 )
 from homeassistant.core import HomeAssistant
-from PIL import Image, ImageOps
+from PIL import Image, ImageOps, UnidentifiedImageError
 
 from .const import (
     AMAZON_DELIVERED,
@@ -201,7 +201,7 @@ async def process_emails(hass: HomeAssistant, config: ConfigEntry) -> dict:
 
     # Copy image file to www directory if enabled
     if config.get(CONF_ALLOW_EXTERNAL):
-        await hass.async_add_executor_job(copy_images, hass, config)
+        await copy_images(hass, config)
 
     return data
 
@@ -226,10 +226,10 @@ def copy_images(hass: HomeAssistant, config: ConfigEntry) -> None:
             except OSError as err:
                 _LOGGER.error("Problem creating: %s, error returned: %s", path, err)
                 return
-        cleanup_images(path)
+        cleanup_images(hass, path)
 
     try:
-        copytree(src, dst, dirs_exist_ok=True)
+        hass.async_add_executor_job(copytree, src, dst, dirs_exist_ok=True)
     except Exception as err:
         _LOGGER.error(
             "Problem copying files from %s to %s error returned: %s", src, dst, err
@@ -367,8 +367,8 @@ async def fetch(
     count = {}
 
     if sensor == "usps_mail":
-        count[sensor] = await hass.async_add_executor_job(
-            get_mails,
+        count[sensor] = get_mails(
+            hass,
             account,
             img_out_path,
             gif_duration,
@@ -377,15 +377,13 @@ async def fetch(
             nomail,
         )
     elif sensor == AMAZON_PACKAGES:
-        count[sensor] = await hass.async_add_executor_job(
-            get_items,
+        count[sensor] = await get_items(
             account,
             ATTR_COUNT,
             amazon_fwds,
             amazon_days,
         )
-        count[AMAZON_ORDER] = await hass.async_add_executor_job(
-            get_items,
+        count[AMAZON_ORDER] = await get_items(
             account,
             ATTR_ORDER,
             amazon_fwds,
@@ -407,7 +405,7 @@ async def fetch(
     elif "_delivering" in sensor:
         prefix = sensor.replace("_delivering", "")
         delivered = await fetch(hass, config, account, data, f"{prefix}_delivered")
-        info = get_count(account, sensor, True)
+        info = await get_count(account, sensor, True)
         count[sensor] = max(0, info[ATTR_COUNT] - delivered)
         count[f"{prefix}_tracking"] = info[ATTR_TRACKING]
     elif sensor == "zpackages_delivered":
@@ -426,9 +424,10 @@ async def fetch(
     elif sensor == "mail_updated":
         count[sensor] = update_time()
     else:
-        count[sensor] = get_count(
+        the_count = await get_count(
             account, sensor, False, img_out_path, hass, amazon_image_name
-        )[ATTR_COUNT]
+        )
+        count[sensor] = the_count[ATTR_COUNT]
 
     data.update(count)
     _LOGGER.debug("Sensor: %s Count: %s", sensor, str(count[sensor]))
@@ -606,6 +605,7 @@ def email_fetch(
 
 
 def get_mails(
+    hass: HomeAssistant,
     account: Type[imaplib.IMAP4_SSL],
     image_output_path: str,
     gif_duration: int,
@@ -636,17 +636,17 @@ def get_mails(
     # Check to see if the path exists, if not make it
     if not os.path.isdir(image_output_path):
         try:
-            os.makedirs(image_output_path)
+            hass.async_add_executor_job(os.makedirs, image_output_path)
         except Exception as err:
             _LOGGER.critical("Error creating directory: %s", str(err))
 
     # Clean up image directory
     _LOGGER.debug("Cleaning up image directory: %s", str(image_output_path))
-    cleanup_images(image_output_path)
+    cleanup_images(hass, image_output_path)
 
     # Copy overlays to image directory
     _LOGGER.debug("Checking for overlay files in: %s", str(image_output_path))
-    copy_overlays(image_output_path)
+    copy_overlays(hass, image_output_path)
 
     if server_response == "OK":
         _LOGGER.debug("Informed Delivery email found processing...")
@@ -677,12 +677,12 @@ def get_mails(
                                 filename = random_filename()
                                 data = str(image["src"]).split(",")[1]
                                 try:
-                                    with open(
-                                        image_output_path + filename, "wb"
-                                    ) as the_file:
-                                        the_file.write(base64.b64decode(data))
-                                        images.append(image_output_path + filename)
-                                        image_count = image_count + 1
+                                    the_file = hass.async_add_executor_job(
+                                        open, image_output_path + filename, "wb"
+                                    )
+                                    the_file.write(base64.b64decode(data))
+                                    images.append(image_output_path + filename)
+                                    image_count = image_count + 1
                                 except Exception as err:
                                     _LOGGER.critical(
                                         "Error opening filepath: %s", str(err)
@@ -699,12 +699,12 @@ def get_mails(
                                 _LOGGER.debug("Discarding junk mail.")
                                 continue
                             try:
-                                with open(
-                                    image_output_path + filename, "wb"
-                                ) as the_file:
-                                    the_file.write(part.get_payload(decode=True))
-                                    images.append(image_output_path + filename)
-                                    image_count = image_count + 1
+                                the_file = hass.async_add_executor_job(
+                                    open, image_output_path + filename, "wb"
+                                )
+                                the_file.write(part.get_payload(decode=True))
+                                images.append(image_output_path + filename)
+                                image_count = image_count + 1
                             except Exception as err:
                                 _LOGGER.critical("Error opening filepath: %s", str(err))
                                 return image_count
@@ -743,51 +743,56 @@ def get_mails(
 
             _LOGGER.debug("Resizing images to 724x320...")
             # Resize images to 724x320
-            all_images = resize_images(images, 724, 320)
+            all_images = hass.async_add_executor_job(
+                resize_images, hass, images, 724, 320
+            )
 
             # Create copy of image list for deleting temporary images
             for image in all_images:
                 images_delete.append(image)
 
             # Create numpy array of images
             _LOGGER.debug("Creating array of image files...")
-            img, *imgs = [Image.open(file) for file in all_images]
+            img, imgs = hass.async_add_executor_job(generate_list, all_images)
 
             try:
                 _LOGGER.debug("Generating animated GIF")
                 # Use Pillow to create mail images
-                img.save(
-                    fp=os.path.join(image_output_path, image_name),
-                    format="GIF",
-                    append_images=imgs,
-                    save_all=True,
-                    duration=gif_duration * 1000,
-                    loop=0,
+                hass.async_add_executor_job(
+                    save_image,
+                    img,
+                    os.path.join(image_output_path, image_name),
+                    imgs,
+                    gif_duration,
                 )
                 _LOGGER.debug("Mail image generated.")
             except Exception as err:
                 _LOGGER.error("Error attempting to generate image: %s", str(err))
             for image in images_delete:
-                cleanup_images(f"{os.path.split(image)[0]}/", os.path.split(image)[1])
+                cleanup_images(
+                    hass, f"{os.path.split(image)[0]}/", os.path.split(image)[1]
+                )
 
         elif image_count == 0:
             _LOGGER.debug("No mail found.")
             if os.path.isfile(image_output_path + image_name):
                 _LOGGER.debug("Removing " + image_output_path + image_name)
-                cleanup_images(image_output_path, image_name)
+                cleanup_images(hass, image_output_path, image_name)
 
             try:
                 _LOGGER.debug("Copying nomail gif")
                 if custom_img is not None:
                     nomail = custom_img
                 else:
                     nomail = os.path.dirname(__file__) + "/mail_none.gif"
-                copyfile(nomail, image_output_path + image_name)
+                hass.async_add_executor_job(
+                    copyfile, nomail, image_output_path + image_name
+                )
             except Exception as err:
                 _LOGGER.error("Error attempting to copy image: %s", str(err))
 
         if gen_mp4:
-            _generate_mp4(image_output_path, image_name)
+            _generate_mp4(hass, image_output_path, image_name)
 
     return image_count
 
@@ -797,7 +802,25 @@ def random_filename(ext: str = ".jpg") -> str:
     return f"{str(uuid.uuid4())}{ext}"
 
 
-def _generate_mp4(path: str, image_file: str) -> None:
+def generate_list(images: list) -> tuple:
+    """Generate list of images."""
+    img, *imgs = [Image.open(file) for file in images]
+    return img, imgs
+
+
+def save_image(img: Image, fp, imgs, duration) -> None:
+    """Save an image."""
+    img.save(
+        fp=fp,
+        format="GIF",
+        append_images=imgs,
+        save_all=True,
+        duration=duration * 1000,
+        loop=0,
+    )
+
+
+def _generate_mp4(hass: HomeAssistant, path: str, image_file: str) -> None:
     """Generate mp4 from gif.
 
     use a subprocess so we don't lock up the thread
@@ -808,7 +831,7 @@ def _generate_mp4(path: str, image_file: str) -> None:
     filecheck = os.path.isfile(mp4_file)
     _LOGGER.debug("Generating mp4: %s", mp4_file)
     if filecheck:
-        cleanup_images(os.path.split(mp4_file))
+        cleanup_images(hass, os.path.split(mp4_file))
         _LOGGER.debug("Removing old mp4: %s", mp4_file)
 
     # TODO: find a way to call ffmpeg the right way from HA
@@ -826,7 +849,7 @@ def _generate_mp4(path: str, image_file: str) -> None:
     )
 
 
-def resize_images(images: list, width: int, height: int) -> list:
+def resize_images(hass: HomeAssistant, images: list, width: int, height: int) -> list:
     """Resize images.
 
     This should keep the aspect ratio of the images
@@ -852,6 +875,13 @@ def resize_images(images: list, width: int, height: int) -> list:
                     img.save(image, img.format)
                     fd_img.close()
                     all_images.append(image)
+                except UnidentifiedImageError as err:
+                    _LOGGER.warning(
+                        "Unable to read image file: %s\nReason: %s",
+                        str(image),
+                        str(err),
+                    )
+                    continue
                 except Exception as err:
                     _LOGGER.error(
                         "Error attempting to read image %s: %s", str(image), str(err)
@@ -864,7 +894,7 @@ def resize_images(images: list, width: int, height: int) -> list:
     return all_images
 
 
-def copy_overlays(path: str) -> None:
+def copy_overlays(hass: HomeAssistant, path: str) -> None:
     """Copy overlay images to image output path."""
     overlays = OVERLAY
     check = all(item in overlays for item in os.listdir(path))
@@ -873,33 +903,34 @@ def copy_overlays(path: str) -> None:
     if not check:
         for file in overlays:
             _LOGGER.debug("Copying file to: %s", str(path + file))
-            copyfile(
+            hass.async_add_executor_job(
+                copyfile,
                 os.path.dirname(__file__) + "/" + file,
                 path + file,
             )
 
 
-def cleanup_images(path: str, image: Optional[str] = None) -> None:
+def cleanup_images(hass: HomeAssistant, path: str, image: Optional[str] = None) -> None:
     """Clean up image storage directory.
 
     Only supose to delete .gif, .mp4, and .jpg files
     """
     if image is not None:
         try:
-            os.remove(path + image)
+            hass.async_add_executor_job(os.remove, path + image)
         except Exception as err:
             _LOGGER.error("Error attempting to remove image: %s", str(err))
         return
 
     for file in os.listdir(path):
         if file.endswith(".gif") or file.endswith(".mp4") or file.endswith(".jpg"):
             try:
-                os.remove(path + file)
+                hass.async_add_executor_job(os.remove, path + file)
             except Exception as err:
                 _LOGGER.error("Error attempting to remove found image: %s", str(err))
 
 
-def get_count(
+async def get_count(
     account: Type[imaplib.IMAP4_SSL],
     sensor_type: str,
     get_tracking_num: bool = False,
@@ -920,7 +951,9 @@ def get_count(
 
     # Return Amazon delivered info
     if sensor_type == AMAZON_DELIVERED:
-        result[ATTR_COUNT] = amazon_search(account, image_path, hass, amazon_image_name)
+        result[ATTR_COUNT] = await amazon_search(
+            account, image_path, hass, amazon_image_name
+        )
         result[ATTR_TRACKING] = ""
         return result
 
@@ -1071,7 +1104,7 @@ def find_text(sdata: Any, account: Type[imaplib.IMAP4_SSL], search_terms: list)
     return count
 
 
-def amazon_search(
+async def amazon_search(
     account: Type[imaplib.IMAP4_SSL],
     image_path: str,
     hass: HomeAssistant,
@@ -1099,8 +1132,7 @@ def amazon_search(
             if server_response == "OK" and data[0] is not None:
                 count += len(data[0].split())
                 _LOGGER.debug("Amazon delivered email(s) found: %s", count)
-                hass.async_add_executor_job(
-                    get_amazon_image,
+                await get_amazon_image(
                     data[0],
                     account,
                     image_path,
@@ -1111,7 +1143,7 @@ def amazon_search(
     return count
 
 
-def get_amazon_image(
+async def get_amazon_image(
     sdata: Any,
     account: Type[imaplib.IMAP4_SSL],
     image_path: str,
@@ -1325,7 +1357,7 @@ def amazon_date_format(arrive_date: str, lang: str) -> tuple:
     return (arrive_date, "%A, %B %d")
 
 
-def get_items(
+async def get_items(
     account: Type[imaplib.IMAP4_SSL],
     param: str = None,
     fwds: Optional[str] = None,

tests/conftest.py

@@ -956,7 +956,7 @@ def test_invalid_ffmpeg():
 
 
 @pytest.fixture
-def mock_copyfile_exception():
+async def mock_copyfile_exception():
     """Fixture to mock copyfile."""
     with patch("custom_components.mail_and_packages.helpers.copyfile") as mock_copyfile:
         mock_copyfile.side_effect = Exception("File not found")
@@ -1363,9 +1363,10 @@ def aioclient_mock_error():
 @pytest.fixture
 def mock_copytree():
     """Fixture to mock copytree."""
-    with patch("custom_components.mail_and_packages.helpers.copytree") as mock_copytree:
-        mock_copytree.return_value = True
-        yield mock_copytree
+    mock_copytree = mock.AsyncMock(return_value=True)
+    # with patch("custom_components.mail_and_packages.helpers.copytree") as mock_copytree:
+    #     mock_copytree.return_value = True
+    yield mock_copytree
 
 
 @pytest.fixture()