Add subscribe mixin (#3710)

* Add subscribe mixin

* Update src/mixins/subscribe-mixin.ts

Co-Authored-By: Paulus Schoutsen <balloob@gmail.com>

* Update src/mixins/subscribe-mixin.ts

Co-Authored-By: Paulus Schoutsen <balloob@gmail.com>

* Update subscribe-mixin.ts

* Update subscribe-mixin.ts

* Add properties

* Fix

src/components/device/ha-device-picker.ts

@@ -14,6 +14,8 @@ import {
   property,
 } from "lit-element";
 import { UnsubscribeFunc } from "home-assistant-js-websocket";
+import { SubscribeMixin } from "../../../src/mixins/subscribe-mixin";
+
 import { HomeAssistant } from "../../types";
 import { fireEvent } from "../../common/dom/fire_event";
 import {
@@ -23,12 +25,11 @@ import {
 import { compare } from "../../common/string/compare";
 
 @customElement("ha-device-picker")
-class HaDevicePicker extends LitElement {
+class HaDevicePicker extends SubscribeMixin(LitElement) {
   @property() public hass?: HomeAssistant;
   @property() public label?: string;
   @property() public value?: string;
   @property() public devices?: DeviceRegistryEntry[];
-  private _unsubDevices?: UnsubscribeFunc;
 
   private _sortedDevices = memoizeOne((devices?: DeviceRegistryEntry[]) => {
     if (!devices || devices.length === 1) {
@@ -39,22 +40,12 @@ class HaDevicePicker extends LitElement {
     return sorted;
   });
 
-  public connectedCallback() {
-    super.connectedCallback();
-    this._unsubDevices = subscribeDeviceRegistry(
-      this.hass!.connection!,
-      (devices) => {
+  public hassSubscribe(): UnsubscribeFunc[] {
+    return [
+      subscribeDeviceRegistry(this.hass!.connection!, (devices) => {
         this.devices = devices;
-      }
-    );
-  }
-
-  public disconnectedCallback() {
-    super.disconnectedCallback();
-    if (this._unsubDevices) {
-      this._unsubDevices();
-      this._unsubDevices = undefined;
-    }
+      }),
+    ];
   }
 
   protected render(): TemplateResult | void {

src/mixins/subscribe-mixin.ts

@@ -0,0 +1,65 @@
+import {
+  LitElement,
+  Constructor,
+  PropertyValues,
+  PropertyDeclarations,
+} from "lit-element";
+import { UnsubscribeFunc } from "home-assistant-js-websocket";
+
+export interface HassSubscribeElement {
+  hassSubscribe(): UnsubscribeFunc[];
+}
+
+/* tslint:disable-next-line */
+export const SubscribeMixin = <T extends LitElement>(
+  superClass: Constructor<T>
+): Constructor<T & HassSubscribeElement> =>
+  // @ts-ignore
+  class extends superClass {
+    /* tslint:disable-next-line */
+    private __unsubs?: UnsubscribeFunc[];
+
+    static get properties(): PropertyDeclarations {
+      return {
+        hass: {},
+      };
+    }
+
+    public connectedCallback() {
+      super.connectedCallback();
+      this.__checkSubscribed();
+    }
+
+    public disconnectedCallback() {
+      super.disconnectedCallback();
+      if (this.__unsubs) {
+        while (this.__unsubs.length) {
+          this.__unsubs.pop()!();
+        }
+        this.__unsubs = undefined;
+      }
+    }
+
+    protected updated(changedProps: PropertyValues) {
+      super.updated(changedProps);
+      if (changedProps.has("hass")) {
+        this.__checkSubscribed();
+      }
+    }
+
+    protected hassSubscribe(): UnsubscribeFunc[] {
+      super.hassSubscribe();
+      return [];
+    }
+
+    private __checkSubscribed(): void {
+      if (
+        this.__unsubs !== undefined ||
+        !((this as unknown) as Element).isConnected ||
+        super.hass === undefined
+      ) {
+        return;
+      }
+      this.__unsubs = this.hassSubscribe();
+    }
+  };