From 97f64f4f0b827667be72d5b7b19722b207d99ae0 Mon Sep 17 00:00:00 2001
From: TheJulianJES <TheJulianJES@users.noreply.github.com>
Date: Tue, 27 Jun 2023 00:12:09 +0200
Subject: [PATCH] Add ZGP basis to quirk gen

---
 quirk_generator/generate_quirk.py              | 12 ++++++++++--
 quirk_generator/templates/imports_template.txt |  2 +-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/quirk_generator/generate_quirk.py b/quirk_generator/generate_quirk.py
index 6ad0f4a..b21e753 100644
--- a/quirk_generator/generate_quirk.py
+++ b/quirk_generator/generate_quirk.py
@@ -4,7 +4,7 @@
 from collections import defaultdict
 
 from jinja2 import Environment, PackageLoader
-from zigpy.profiles import zha, zll
+from zigpy.profiles import zgp, zha, zll
 from zigpy.zcl import clusters
 
 _LOGGER = logging.getLogger(__name__)
@@ -17,10 +17,11 @@ def process_profiles(data: dict) -> dict:
     """Process the profiles in the device signature."""
     profiles = set()
     for endpoint in data["signature"]["endpoints"].values():
-        profiles.add(endpoint["profile_id"])
+        profiles.add(int(endpoint["profile_id"], 16))
     data["present_profiles"] = profiles
     data["zha_profile"] = zha.PROFILE_ID
     data["zll_profile"] = zll.PROFILE_ID
+    data["zgp_profile"] = zgp.PROFILE_ID
     return data
 
 
@@ -43,6 +44,8 @@ def process_profile_id(profile_id: str) -> str | int:
         return "zha.PROFILE_ID"
     elif profile_id == zll.PROFILE_ID:
         return "zll.PROFILE_ID"
+    elif profile_id == zgp.PROFILE_ID:
+        return "zgp.PROFILE_ID"
     else:
         return profile_id
 
@@ -62,6 +65,11 @@ def process_device_type(profile_id: str, device_type: str) -> str | int:
             return f"zll.DeviceType.{zll.DeviceType(device_type).name}"
         except ValueError:
             return device_type
+    elif profile_id == zgp.PROFILE_ID:
+        try:
+            return f"zgp.DeviceType.{zgp.DeviceType(device_type).name}"
+        except ValueError:
+            return device_type
     else:
         return device_type
 
diff --git a/quirk_generator/templates/imports_template.txt b/quirk_generator/templates/imports_template.txt
index 7067315..f7bf886 100644
--- a/quirk_generator/templates/imports_template.txt
+++ b/quirk_generator/templates/imports_template.txt
@@ -1,4 +1,4 @@
-from zigpy.profiles import {% if zha_profile in present_profiles -%}zha{% if present_profiles |length > 1 -%},{% endif %} {% endif %}{% if zll_profile in present_profiles -%}zll{% endif %}
+from zigpy.profiles import {% if zha_profile in present_profiles -%}zha{% if present_profiles |length > 1 -%},{% endif %} {% endif %}{% if zll_profile in present_profiles -%}zll{% if present_profiles |length > 2 -%},{% endif %}{% endif %}{% if zgp_profile in present_profiles -%}zgp{% endif %}
 from zigpy.quirks import CustomDevice
 {% for module, classes in grouped_clusters.items() %}from {{module}} import (
 {% for class in classes %}    {{class}},{% if not loop.last %}