better working ws

Signed-off-by: Jess Frazelle <github@jessfraz.com>

generate/functions-ws.py.jinja2

@@ -1,4 +1,6 @@
 from typing import Any, Dict, Optional, Union, List
+import json
+import bson
 
 from websockets.sync.client import connect as ws_connect
 from websockets.client import connect as ws_connect_async
@@ -16,26 +18,37 @@ from ...types import Response
 
 def _get_kwargs(
     {% for arg in args %}
+    {% if arg.in_query %}
     {% if arg.is_optional == False %}
     {{arg.name}}: {{arg.type}},
     {% endif %}
+    {% endif %}
     {% endfor %}
     *,
     client: Client,
     {% for arg in args %}
+    {% if arg.in_query %}
     {% if arg.is_optional %}
     {{arg.name}}: {{arg.type}},
     {% endif %}
+    {% endif %}
     {% endfor %}
 ) -> Dict[str, Any]:
     url = "{{url_template}}".format(client.base_url{% for arg in args %}{% if arg.in_url %}, {{arg.name}}={{arg.name}}{% endif %}{% endfor %}) # noqa: E501
     {% for arg in args %}
     {% if arg.in_query %}
     if {{arg.name}} is not None:
+        {% if arg.type == "bool" %}
+        if "?" in url:
+            url = url + "&{{arg.name}}=" + str({{arg.name}}).lower()
+        else:
+            url = url + "?{{arg.name}}=" + str({{arg.name}}).lower()
+        {% else %}
         if "?" in url:
             url = url + "&{{arg.name}}=" + str({{arg.name}})
         else:
             url = url + "?{{arg.name}}=" + str({{arg.name}})
+        {% endif %}
     {% endif %}
     {% endfor %}
 
@@ -48,29 +61,34 @@ def _get_kwargs(
         "headers": headers,
         "cookies": cookies,
         "timeout": client.get_timeout(),
-        {% if has_request_body %}"content": body,{% endif %}
     }
 
 
 def sync(
     {% for arg in args %}
+    {% if arg.in_query %}
     {% if arg.is_optional == False %}
     {{arg.name}}: {{arg.type}},
     {% endif %}
+    {% endif %}
     {% endfor %}
     *,
     client: Client,
     {% for arg in args %}
+    {% if arg.in_query %}
     {% if arg.is_optional %}
     {{arg.name}}: {{arg.type}},
     {% endif %}
+    {% endif %}
     {% endfor %}
 ) -> ClientConnection:
     {%if docs%}"""{{docs}}"""  # noqa: E501{% endif %}
 
     kwargs = _get_kwargs(
         {% for arg in args %}
+        {% if arg.in_query %}
         {{arg.name}}={{arg.name}},
+        {% endif %}
         {% endfor %}
         client=client,
     )
@@ -85,23 +103,29 @@ def sync(
 
 async def asyncio(
     {% for arg in args %}
+    {% if arg.in_query %}
     {% if arg.is_optional == False %}
     {{arg.name}}: {{arg.type}},
     {% endif %}
+    {% endif %}
     {% endfor %}
     *,
     client: Client,
     {% for arg in args %}
+    {% if arg.in_query %}
     {% if arg.is_optional %}
     {{arg.name}}: {{arg.type}},
     {% endif %}
+    {% endif %}
     {% endfor %}
 ) -> WebSocketClientProtocol:
     {%if docs%}"""{{docs}}"""  # noqa: E501{% endif %}
 
     kwargs = _get_kwargs(
         {% for arg in args %}
+        {% if arg.in_query %}
         {{arg.name}}={{arg.name}},
+        {% endif %}
         {% endfor %}
         client=client,
     )
@@ -111,3 +135,48 @@ async def asyncio(
 
     # Return an error if we got here.
     return Error(message="An error occurred while connecting to the websocket.")
+
+{% if has_request_body %}
+class WebSocket:
+    """A websocket connection to the API endpoint."""
+    ws:  ClientConnection
+
+    def __init__(self,
+        {% for arg in args %}
+        {% if arg.in_query %}
+        {% if arg.is_optional == False %}
+        {{arg.name}}: {{arg.type}},
+        {% endif %}
+        {% endif %}
+        {% endfor %}
+        client: Client,
+        {% for arg in args %}
+        {% if arg.in_query %}
+        {% if arg.is_optional %}
+        {{arg.name}}: {{arg.type}},
+        {% endif %}
+        {% endif %}
+        {% endfor %}
+    ):
+        self.ws = sync(
+        {% for arg in args %}
+        {% if arg.in_query %}
+            {{arg.name}},
+        {% endif %}
+        {% endfor %}
+            client=client,
+        )
+
+    def send(self, data:{% for arg in args %}{%if arg.name == "body" %}{{arg.type}}{% endif %}{% endfor %}):
+        """Send data to the websocket."""
+        self.ws.send(json.dumps(data.to_dict()))
+
+    def send_binary(self, data:{% for arg in args %}{%if arg.name == "body" %}{{arg.type}}{% endif %}{% endfor %}):
+        """Send data as bson to the websocket."""
+        self.ws.send(bson.BSON.encode(data.to_dict()))
+
+    def recv(self) ->  {{response_type}}:
+        """Receive data from the websocket."""
+        message = self.ws.recv()
+        return {{response_type}}.from_dict(json.loads(message))
+{%endif%}