Add functionality to bulk rename friendly names

homeassistant-entity-renamer.py

@@ -1,12 +1,12 @@
 #!/usr/bin/env python3
 
 import argparse
+import tabulate
 import config
 import csv
 import json
 import re
 import requests
-import tabulate
 import websocket
 
 tabulate.PRESERVE_WHITESPACE = True
@@ -25,7 +25,7 @@ def align_strings(table):
 
     if len(table) == 0:
         return
-    
+
     for column in range(len(table[0])):
         # Get the column data from the table
         column_data = [row[column] for row in table]
@@ -34,7 +34,7 @@ def align_strings(table):
         strings_to_align = [s for s in column_data if alignment_char in s]
         if len(strings_to_align) == 0:
             continue
-        
+
         max_length = max([len(s.split(alignment_char)[0]) for s in strings_to_align])
 
         def align_string(s):
@@ -108,14 +108,14 @@ def process_entities(entity_data, search_regex, replace_regex=None, output_csv=N
     # Ask user for confirmation if replace_regex is provided
     if not replace_regex:
         return
-    
+
     answer = input("\nDo you want to proceed with renaming the entities? (y/N): ")
     if answer.lower() not in ["y", "yes"]:
         print("Renaming process aborted.")
         return
 
     rename_entities(rename_data)
-    
+
 def rename_entities(rename_data):
     websocket_url = f'ws{TLS_S}://{config.HOST}/api/websocket'
     ws = websocket.WebSocket()
@@ -151,19 +151,151 @@ def rename_entities(rename_data):
     ws.close()
 
 
+def list_friendly_names(regex=None):
+    # API endpoint for retrieving all entities
+    api_endpoint = f"http{TLS_S}://{config.HOST}/api/states"
+
+    # Send GET request to the API endpoint
+    response = requests.get(api_endpoint, headers=headers)
+
+    # Check if the request was successful
+    if response.status_code == 200:
+        # Parse the JSON response
+        data = json.loads(response.text)
+
+        # Extract entity IDs and friendly names
+        entity_data = [
+            (entity["attributes"].get("friendly_name", ""), entity["entity_id"])
+            for entity in data
+        ]
+
+        # Filter the entity data if regex argument is provided
+        if regex:
+            filtered_entity_data = [
+                (friendly_name, entity_id)
+                for friendly_name, entity_id in entity_data
+                if re.search(regex, friendly_name)
+            ]
+            entity_data = filtered_entity_data
+
+        # Sort the entity data by friendly name
+        entity_data = sorted(entity_data, key=lambda x: x[0])
+
+        # Output the entity data
+        return entity_data
+
+    else:
+        print(f"Error: {response.status_code} - {response.text}")
+        return []
+
+
+def process_friendly_names(entity_data, search_regex, replace_regex=None):
+    rename_data = []
+    if replace_regex:
+        for friendly_name, entity_id in entity_data:
+            new_friendly_name = re.sub(search_regex, replace_regex, friendly_name)
+            rename_data.append((friendly_name, entity_id, new_friendly_name))
+    else:
+        rename_data = [
+            (friendly_name, entity_id, "") for friendly_name, entity_id in entity_data
+        ]
+
+    # Print the table with friendly name and entity ID
+    table = [
+        ("Friendly Name", "Current Entity ID", "New Friendly Name")
+    ] + align_strings(rename_data)
+    print(tabulate.tabulate(table, headers="firstrow", tablefmt="github"))
+
+    # Ask user for confirmation if replace_regex is provided
+    if not replace_regex:
+        return
+
+    answer = input(
+        "\nDo you want to proceed with renaming the entities' friendly names? (y/N): "
+    )
+    if answer.lower() not in ["y", "yes"]:
+        print("Renaming process aborted.")
+        return
+
+    rename_friendly_names(rename_data)
+
+
+def rename_friendly_names(rename_data):
+    websocket_url = f"ws{TLS_S}://{config.HOST}/api/websocket"
+    ws = websocket.WebSocket()
+    ws.connect(websocket_url)
+
+    auth_req = ws.recv()
+
+    # Authenticate with Home Assistant
+    auth_msg = json.dumps({"type": "auth", "access_token": config.ACCESS_TOKEN})
+    ws.send(auth_msg)
+    auth_result = ws.recv()
+    auth_result = json.loads(auth_result)
+    if auth_result["type"] != "auth_ok":
+        print("Authentication failed. Check your access token.")
+        return
+
+    # Rename the entities
+    for index, (_, entity_id, friendly_name) in enumerate(rename_data, start=1):
+        entity_registry_update_msg = json.dumps(
+            {
+                "id": index,
+                "type": "config/entity_registry/update",
+                "entity_id": entity_id,
+                "new_entity_id": entity_id,
+                "name": friendly_name,
+            }
+        )
+        ws.send(entity_registry_update_msg)
+        update_result = ws.recv()
+        update_result = json.loads(update_result)
+        if update_result["success"]:
+            print(f"Entity '{entity_id}' renamed to '{friendly_name}' successfully!")
+        else:
+            print(
+                f"Failed to rename entity '{entity_id}': {update_result['error']['message']}"
+            )
+
+    ws.close()
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="HomeAssistant Entity Renamer")
-    parser.add_argument('--search', dest='search_regex', help='Regular expression for search. Note: Only searches entity IDs.')
-    parser.add_argument('--replace', dest='replace_regex', help='Regular expression for replace')
-    parser.add_argument('--output-csv', dest='output_csv', help='Output preview table to CSV.')
+    parser.add_argument(
+        "-s",
+        "--search",
+        dest="search_regex",
+        help="Regular expression for search. Note: Only searches entity IDs. Use -f to search friendly names.",
+    )
+    parser.add_argument(
+        "-r", "--replace", dest="replace_regex", help="Regular expression for replace"
+    )
+    parser.add_argument(
+        "-f",
+        "--rename-friendly-name",
+        dest="rename_friendly_name",
+        action="store_true",
+        help="Search for and rename the friendly names of the entities.",
+    )
+    parser.add_argument(
+        "-o", "--output-csv", dest="output_csv", help="Output preview table to CSV."
+    )
     args = parser.parse_args()
 
-    if args.search_regex:
+    if not args.rename_friendly_name and args.search_regex:
         entity_data = list_entities(args.search_regex)
 
         if entity_data:
             process_entities(entity_data, args.search_regex, args.replace_regex, args.output_csv)
         else:
             print("No entities found matching the search regex.")
+    elif args.search_regex:
+        entity_data = list_friendly_names(args.search_regex)
+
+        if entity_data:
+            process_friendly_names(entity_data, args.search_regex, args.replace_regex)
+        else:
+            print("No names found matching the search regex.")
     else:
         parser.print_help()