From 348b7abe7d7d2c6f792a3ab7139cfd85876180d0 Mon Sep 17 00:00:00 2001
From: Flyte <flyte@failcode.co.uk>
Date: Wed, 17 Feb 2016 18:12:36 +0000
Subject: [PATCH] Change TCP component to use Jinja2 instead of regex

---
 homeassistant/components/binary_sensor/tcp.py |  2 +-
 homeassistant/components/sensor/tcp.py        | 41 +++++++++++--------
 homeassistant/components/tcp.py               |  1 +
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/binary_sensor/tcp.py b/homeassistant/components/binary_sensor/tcp.py
index 056a59fbdf1..e6d01b893df 100644
--- a/homeassistant/components/binary_sensor/tcp.py
+++ b/homeassistant/components/binary_sensor/tcp.py
@@ -19,7 +19,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
     """ Create the BinarySensor. """
     if not BinarySensor.validate_config(config):
         return False
-    add_entities((BinarySensor(config),))
+    add_entities((BinarySensor(hass, config),))
 
 
 class BinarySensor(Sensor, BinarySensorDevice):
diff --git a/homeassistant/components/sensor/tcp.py b/homeassistant/components/sensor/tcp.py
index b45f35434af..29d6c260fdd 100644
--- a/homeassistant/components/sensor/tcp.py
+++ b/homeassistant/components/sensor/tcp.py
@@ -5,14 +5,16 @@ Provides a sensor which gets its values from a TCP socket.
 """
 import logging
 import socket
-import re
 from select import select
 
 from homeassistant.const import CONF_NAME, CONF_HOST
+from homeassistant.util import template
+from homeassistant.exceptions import TemplateError
 from homeassistant.helpers.entity import Entity
 from homeassistant.components.tcp import (
     DOMAIN, CONF_PORT, CONF_TIMEOUT, CONF_PAYLOAD, CONF_UNIT, CONF_VALUE_REGEX,
-    CONF_VALUE_ON, CONF_BUFFER_SIZE, DEFAULT_TIMEOUT, DEFAULT_BUFFER_SIZE
+    CONF_VALUE_TEMPLATE, CONF_VALUE_ON, CONF_BUFFER_SIZE, DEFAULT_TIMEOUT,
+    DEFAULT_BUFFER_SIZE
 )
 
 
@@ -25,15 +27,16 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
     """ Create the Sensor. """
     if not Sensor.validate_config(config):
         return False
-    add_entities((Sensor(config),))
+    add_entities((Sensor(hass, config),))
 
 
 class Sensor(Entity):
     """ Sensor Entity which gets its value from a TCP socket. """
     required = tuple()
 
-    def __init__(self, config):
+    def __init__(self, hass, config):
         """ Set all the config values if they exist and get initial state. """
+        self._hass = hass
         self._config = {
             CONF_NAME: config.get(CONF_NAME),
             CONF_HOST: config[CONF_HOST],
@@ -42,6 +45,7 @@ class Sensor(Entity):
             CONF_PAYLOAD: config[CONF_PAYLOAD],
             CONF_UNIT: config.get(CONF_UNIT),
             CONF_VALUE_REGEX: config.get(CONF_VALUE_REGEX),
+            CONF_VALUE_TEMPLATE: config.get(CONF_VALUE_TEMPLATE),
             CONF_VALUE_ON: config.get(CONF_VALUE_ON),
             CONF_BUFFER_SIZE: config.get(
                 CONF_BUFFER_SIZE, DEFAULT_BUFFER_SIZE),
@@ -78,6 +82,7 @@ class Sensor(Entity):
     def update(self):
         """ Get the latest value for this sensor. """
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
         try:
             sock.connect((self._config[CONF_HOST], self._config[CONF_PORT]))
         except socket.error as err:
@@ -85,6 +90,7 @@ class Sensor(Entity):
                 "Unable to connect to %s on port %s: %s",
                 self._config[CONF_HOST], self._config[CONF_PORT], err)
             return
+
         try:
             sock.send(self._config[CONF_PAYLOAD].encode())
         except socket.error as err:
@@ -93,6 +99,7 @@ class Sensor(Entity):
                 self._config[CONF_PAYLOAD], self._config[CONF_HOST],
                 self._config[CONF_PORT], err)
             return
+
         readable, _, _ = select([sock], [], [], self._config[CONF_TIMEOUT])
         if not readable:
             _LOGGER.warning(
@@ -101,20 +108,20 @@ class Sensor(Entity):
                 self._config[CONF_TIMEOUT], self._config[CONF_PAYLOAD],
                 self._config[CONF_HOST], self._config[CONF_PORT])
             return
+
         value = sock.recv(self._config[CONF_BUFFER_SIZE]).decode()
-        if self._config[CONF_VALUE_REGEX] is not None:
-            match = re.match(self._config[CONF_VALUE_REGEX], value)
-            if match is None:
-                _LOGGER.warning(
-                    "Unable to match value using value_regex of %r: %r",
-                    self._config[CONF_VALUE_REGEX], value)
-                return
+
+        if self._config[CONF_VALUE_TEMPLATE] is not None:
             try:
-                self._state = match.groups()[0]
-            except IndexError:
-                _LOGGER.error(
-                    "You must include a capture group in the regex for %r: %r",
-                    self.name, self._config[CONF_VALUE_REGEX])
+                self._state = template.render(
+                    self._hass,
+                    self._config[CONF_VALUE_TEMPLATE],
+                    value=value)
                 return
-            return
+            except TemplateError as err:
+                _LOGGER.error(
+                    "Unable to render template of %r with value: %r",
+                    self._config[CONF_VALUE_TEMPLATE], value)
+                return
+
         self._state = value
diff --git a/homeassistant/components/tcp.py b/homeassistant/components/tcp.py
index 9b09ceaf68e..c82bf935910 100644
--- a/homeassistant/components/tcp.py
+++ b/homeassistant/components/tcp.py
@@ -10,6 +10,7 @@ CONF_TIMEOUT = "timeout"
 CONF_PAYLOAD = "payload"
 CONF_UNIT = "unit"
 CONF_VALUE_REGEX = "value_regex"
+CONF_VALUE_TEMPLATE = "value_template"
 CONF_VALUE_ON = "value_on"
 CONF_BUFFER_SIZE = "buffer_size"