Add YAML output

twhatter/cli.py

@@ -5,7 +5,7 @@
 import click
 import IPython
 
-from twhatter.output import Print, Json, Database
+from twhatter.output import Print, Json, Database, Yaml
 from twhatter.output.sqlalchemy import Tweet, User
 from twhatter.log import log_setup
 
@@ -34,6 +34,14 @@ def db(ctx, db_url):
 def json(ctx, json_file):
     ctx.obj['output'] = Json(json_file)
 
+
+@main.group()
+@click.option('-f', '--yaml_file', type=str, default="/tmp/output.yaml", show_default=True)
+@click.pass_context
+def yaml(ctx, yaml_file):
+    ctx.obj['output'] = Yaml(yaml_file)
+
+
 @main.command()
 @click.option('-l', '--limit', type=int, default=100, show_default=True)
 @click.argument('user')
@@ -70,6 +78,8 @@ db.add_command(timeline)
 json.add_command(profile)
 json.add_command(timeline)
 
+yaml.add_command(profile)
+yaml.add_command(timeline)
 
 if __name__ == "__main__":
     main(obj={})

twhatter/output/__init__.py

@@ -2,3 +2,4 @@ from .print import Print
 from .base import OutputBase
 from .json import Json
 from .sqlalchemy import Database
+from .yaml import Yaml

twhatter/output/yaml.py

@@ -0,0 +1,34 @@
+import yaml
+import logging
+from bs4 import PageElement
+
+from .base import OutputBase
+from twhatter.client import ClientTimeline, ClientProfile
+
+
+logger = logging.getLogger(__name__)
+
+
+def PageElement_representer(dumper, data):
+    return dumper.represent_data(None)
+
+
+yaml.add_multi_representer(PageElement, PageElement_representer)
+
+
+class Yaml(OutputBase):
+    def __init__(self, yaml_path):
+        logger.info("Output set to {}".format(yaml_path))
+        self.yaml_path = yaml_path
+
+    def output_tweets(self, user, limit):
+        client_timeline = ClientTimeline(user, limit)
+
+        with open(self.yaml_path, 'w') as f:
+            yaml.dump([t for t in client_timeline], f, indent=2)
+
+    def output_user(self, user):
+        p = ClientProfile(user)
+
+        with open(self.yaml_path, 'w') as f:
+            yaml.dump(p.user, f, indent=2, default_flow_style=False)