<commit_before>import logging
from importlib import import_module

import peewee as pw
from cached_property import cached_property
from flask._compat import with_metaclass, string_types
from peewee_migrate.router import Router, LOGGER
from playhouse.db_url import connect

__license__ = "MIT"
__project__ = "Flask-PW"
__version__ = "0.0.2"

class Choices:

    """Model's choices helper."""

    def __init__(self, *choices):
        """Parse provided choices."""
        self._choices = []
        self._reversed = {}
        for choice in choices:
            if isinstance(choice, string_types):
                choice = (choice, choice)
            self._choices.append(choice)
            self._reversed[str(choice[1])] = choice[0]

    def __getattr__(self, name, default=None):
        """Get choice value by name."""
        return self._reversed.get(name, default)

    def __iter__(self):
        """Iterate self."""
        return iter(self._choices)

    def __str__(self):
        """String representation."""
        return ", ".join(self._reversed.keys())

    def __repr__(self):
        """Python representation."""
        return "<Choices: %s>" % self

    def __nonzero__(self):
        return True

class Signal:

    """Simplest signals implementation.

    ::
        @Model.post_save
        def example(instance, created=False):
            pass

    """

    __slots__ = 'receivers'

    def __init__(self):
        """Initialize the signal."""
        self.receivers = []

    def connect(self, receiver):
        """Append receiver."""
        if not callable(receiver):
            raise ValueError('Invalid receiver: %s' % receiver)
        self.receivers.append(receiver)

    def __call__(self, receiver):
        """Support decorators.."""
        self.connect(receiver)
        return receiver

    def disconnect(self, receiver):
        """Remove receiver."""
        try:
            self.receivers.remove(receiver)
        except ValueError:
            raise ValueError('Unknown receiver: %s' % receiver)

    def send(self, instance, *args, **kwargs):
        """Send signal."""
        for receiver in self.receivers:
            receiver(instance, *args, **kwargs)

class BaseSignalModel(pw.BaseModel):

    """Create signals."""

    models = []

    def __new__(mcs, name, bases, attrs):
        """Create signals."""
        cls = super(BaseSignalModel, mcs).__new__(mcs, name, bases, attrs)
        cls.pre_save = Signal()
        cls.pre_delete = Signal()
        cls.post_delete = Signal()
        cls.post_save = Signal()

        if cls._meta.db_table and cls._meta.db_table != 'model':
            mcs.models.append(cls)
        return cls

class Model(with_metaclass(BaseSignalModel, pw.Model)):

    @property
    def pk(self):
        """Return primary key value."""
        return self._get_pk_value()

    @classmethod
    def get_or_none(cls, *args, **kwargs):
        try:
            return cls.get(*args, **kwargs)
        except cls.DoesNotExist:
            return None

    def save(self, force_insert=False, **kwargs):
        """Send signals."""
        created = force_insert or not bool(self.pk)
        self.pre_save.send(self, created=created)
        super(Model, self).save(force_insert=force_insert, **kwargs)
        self.post_save.send(self, created=created)

    def delete_instance(self, *args, **kwargs):
        """Send signals."""
        self.pre_delete.send(self)
        super(Model, self).delete_instance(*args, **kwargs)
        self.post_delete.send(self)

class Peewee(object):

    def __init__(self, app=None):
        """Initialize the plugin."""
        self.app = app
        self.database = pw.Proxy()
        if app is not None:
            self.init_app(app)

    def init_app(self, app, database=None):
        """Initialize application."""

        # Register application
        if not app:
            raise RuntimeError('Invalid application.')
        self.app = app
        if not hasattr(app, 'extensions'):
            app.extensions = {}
        app.extensions['peewee'] = self

        app.config.setdefault('PEEWEE_CONNECTION_PARAMS', {})
        app.config.setdefault('PEEWEE_DATABASE_URI', 'sqlite:///peewee.sqlite')
        app.config.setdefault('PEEWEE_MODELS_IGNORE', [])
        app.config.setdefault('PEEWEE_MANUAL', False)
        app.config.setdefault('PEEWEE_MIGRATIONS', 'migrations')
        app.config.setdefault('PEEWEE_MODELS_CLASS', Model)
        app.config.setdefault('PEEWEE_MODELS_MODULE', '')

        # Initialize database
        database = database or app.config.get('PEEWEE_DATABASE_URI')
        if not database:
            raise RuntimeError('Invalid database.')
        if isinstance(database, string_types):
            database = connect(database, **app.config['PEEWEE_CONNECTION_PARAMS'])
        self.database.initialize(database)
        if self.database.database == ':memory:':
            app.config['PEEWEE_MANUAL'] = True
        if not app.config['PEEWEE_MANUAL']:
            app.before_request(self.connect)
            app.after_request(self.close)

    def connect(self):
        """Initialize connection to databse."""
        return self.database.connect()

    def close(self, response):
        """Close connection to database."""
        if not self.database.is_closed():
            self.database.close()
        return response

    @cached_property
    def Model(self):
        """Bind model to self database."""
        Model_ = self.app.config['PEEWEE_MODELS_CLASS']
        Meta = type('Meta', (), {'database': self.database})
        return type('Model', (Model_,), {'Meta': Meta})

    @property
    def models(self):
        """Return self.application models."""
        Model_ = self.app.config['PEEWEE_MODELS_CLASS']
        ignore = self.app.config['PEEWEE_MODELS_IGNORE']

        models = []
        if Model_ is not Model:
            try:
                mod = import_module(self.app.config['PEEWEE_MODELS_MODULE'])
                for model in dir(mod):
                    models = getattr(mod, model)
                    if not isinstance(model, pw.Model):
                        continue
                    models.append(models)
            except ImportError:
                return models
        elif isinstance(Model_, BaseSignalModel):
            models = BaseSignalModel.models

        return [m for m in models if m._meta.name not in ignore]

    @cached_property
    def manager(self):
        """Integrate a Flask-Script."""
        from flask_script import Manager

        manager = Manager(usage="Migrate database.")

        @manager.command
        def create(name, auto=False):
            """Create a new migration."""

            LOGGER.setLevel('INFO')
            LOGGER.propagate = 0

            router = Router(self.database, self.app.config['PEEWEE_MIGRATIONS'])

            if auto:
                auto = self.models

            router.create(name, auto=auto)

        @manager.command
        def migrate(name=None, fake=False):
            """Run migrations."""
            from peewee_migrate.router import Router, LOGGER

            LOGGER.setLevel('INFO')
            LOGGER.propagate = 0

            router = Router(self.database, self.app.config['PEEWEE_MIGRATIONS'])
            migrations = router.run(name, fake=fake)
            if migrations:
                logging.warn('Migrations are completed: %s' % ', '.join(migrations))

        @manager.command
        def rollback(name):
            """Rollback migrations."""
            from peewee_migrate.router import Router, LOGGER

            LOGGER.setLevel('INFO')
            LOGGER.propagate = 0

            router = Router(self.database, self.app.config['PEEWEE_MIGRATIONS'])
            router.rollback(name)

        @manager.command
        def list():
            """List migrations."""
            from peewee_migrate.router import Router, LOGGER

            LOGGER.setLevel('DEBUG')
            LOGGER.propagate = 0

            router = Router(self.database, self.app.config['PEEWEE_MIGRATIONS'])
            LOGGER.info('Migrations are done:')
            LOGGER.info('\n'.join(router.done))
            LOGGER.info()
            LOGGER.info('Migrations are undone:')
            LOGGER.info('\n'.join(router.diff))

        return manager
<commit_msg>Fix list<commit_after>import logging
from importlib import import_module

import peewee as pw
from cached_property import cached_property
from flask._compat import with_metaclass, string_types
from peewee_migrate.router import Router, LOGGER
from playhouse.db_url import connect

__license__ = "MIT"
__project__ = "Flask-PW"
__version__ = "0.0.2"

class Choices:

    """Model's choices helper."""

    def __init__(self, *choices):
        """Parse provided choices."""
        self._choices = []
        self._reversed = {}
        for choice in choices:
            if isinstance(choice, string_types):
                choice = (choice, choice)
            self._choices.append(choice)
            self._reversed[str(choice[1])] = choice[0]

    def __getattr__(self, name, default=None):
        """Get choice value by name."""
        return self._reversed.get(name, default)

    def __iter__(self):
        """Iterate self."""
        return iter(self._choices)

    def __str__(self):
        """String representation."""
        return ", ".join(self._reversed.keys())

    def __repr__(self):
        """Python representation."""
        return "<Choices: %s>" % self

    def __nonzero__(self):
        return True

class Signal:

    """Simplest signals implementation.

    ::
        @Model.post_save
        def example(instance, created=False):
            pass

    """

    __slots__ = 'receivers'

    def __init__(self):
        """Initialize the signal."""
        self.receivers = []

    def connect(self, receiver):
        """Append receiver."""
        if not callable(receiver):
            raise ValueError('Invalid receiver: %s' % receiver)
        self.receivers.append(receiver)

    def __call__(self, receiver):
        """Support decorators.."""
        self.connect(receiver)
        return receiver

    def disconnect(self, receiver):
        """Remove receiver."""
        try:
            self.receivers.remove(receiver)
        except ValueError:
            raise ValueError('Unknown receiver: %s' % receiver)

    def send(self, instance, *args, **kwargs):
        """Send signal."""
        for receiver in self.receivers:
            receiver(instance, *args, **kwargs)

class BaseSignalModel(pw.BaseModel):

    """Create signals."""

    models = []

    def __new__(mcs, name, bases, attrs):
        """Create signals."""
        cls = super(BaseSignalModel, mcs).__new__(mcs, name, bases, attrs)
        cls.pre_save = Signal()
        cls.pre_delete = Signal()
        cls.post_delete = Signal()
        cls.post_save = Signal()

        if cls._meta.db_table and cls._meta.db_table != 'model':
            mcs.models.append(cls)
        return cls

class Model(with_metaclass(BaseSignalModel, pw.Model)):

    @property
    def pk(self):
        """Return primary key value."""
        return self._get_pk_value()

    @classmethod
    def get_or_none(cls, *args, **kwargs):
        try:
            return cls.get(*args, **kwargs)
        except cls.DoesNotExist:
            return None

    def save(self, force_insert=False, **kwargs):
        """Send signals."""
        created = force_insert or not bool(self.pk)
        self.pre_save.send(self, created=created)
        super(Model, self).save(force_insert=force_insert, **kwargs)
        self.post_save.send(self, created=created)

    def delete_instance(self, *args, **kwargs):
        """Send signals."""
        self.pre_delete.send(self)
        super(Model, self).delete_instance(*args, **kwargs)
        self.post_delete.send(self)

class Peewee(object):

    def __init__(self, app=None):
        """Initialize the plugin."""
        self.app = app
        self.database = pw.Proxy()
        if app is not None:
            self.init_app(app)

    def init_app(self, app, database=None):
        """Initialize application."""

        # Register application
        if not app:
            raise RuntimeError('Invalid application.')
        self.app = app
        if not hasattr(app, 'extensions'):
            app.extensions = {}
        app.extensions['peewee'] = self

        app.config.setdefault('PEEWEE_CONNECTION_PARAMS', {})
        app.config.setdefault('PEEWEE_DATABASE_URI', 'sqlite:///peewee.sqlite')
        app.config.setdefault('PEEWEE_MODELS_IGNORE', [])
        app.config.setdefault('PEEWEE_MANUAL', False)
        app.config.setdefault('PEEWEE_MIGRATIONS', 'migrations')
        app.config.setdefault('PEEWEE_MODELS_CLASS', Model)
        app.config.setdefault('PEEWEE_MODELS_MODULE', '')

        # Initialize database
        database = database or app.config.get('PEEWEE_DATABASE_URI')
        if not database:
            raise RuntimeError('Invalid database.')
        if isinstance(database, string_types):
            database = connect(database, **app.config['PEEWEE_CONNECTION_PARAMS'])
        self.database.initialize(database)
        if self.database.database == ':memory:':
            app.config['PEEWEE_MANUAL'] = True
        if not app.config['PEEWEE_MANUAL']:
            app.before_request(self.connect)
            app.after_request(self.close)

    def connect(self):
        """Initialize connection to databse."""
        return self.database.connect()

    def close(self, response):
        """Close connection to database."""
        if not self.database.is_closed():
            self.database.close()
        return response

    @cached_property
    def Model(self):
        """Bind model to self database."""
        Model_ = self.app.config['PEEWEE_MODELS_CLASS']
        Meta = type('Meta', (), {'database': self.database})
        return type('Model', (Model_,), {'Meta': Meta})

    @property
    def models(self):
        """Return self.application models."""
        Model_ = self.app.config['PEEWEE_MODELS_CLASS']
        ignore = self.app.config['PEEWEE_MODELS_IGNORE']

        models = []
        if Model_ is not Model:
            try:
                mod = import_module(self.app.config['PEEWEE_MODELS_MODULE'])
                for model in dir(mod):
                    models = getattr(mod, model)
                    if not isinstance(model, pw.Model):
                        continue
                    models.append(models)
            except ImportError:
                return models
        elif isinstance(Model_, BaseSignalModel):
            models = BaseSignalModel.models

        return [m for m in models if m._meta.name not in ignore]

    @cached_property
    def manager(self):
        """Integrate a Flask-Script."""
        from flask_script import Manager

        manager = Manager(usage="Migrate database.")

        @manager.command
        def create(name, auto=False):
            """Create a new migration."""

            LOGGER.setLevel('INFO')
            LOGGER.propagate = 0

            router = Router(self.database, self.app.config['PEEWEE_MIGRATIONS'])

            if auto:
                auto = self.models

            router.create(name, auto=auto)

        @manager.command
        def migrate(name=None, fake=False):
            """Run migrations."""
            from peewee_migrate.router import Router, LOGGER

            LOGGER.setLevel('INFO')
            LOGGER.propagate = 0

            router = Router(self.database, self.app.config['PEEWEE_MIGRATIONS'])
            migrations = router.run(name, fake=fake)
            if migrations:
                logging.warn('Migrations are completed: %s' % ', '.join(migrations))

        @manager.command
        def rollback(name):
            """Rollback migrations."""
            from peewee_migrate.router import Router, LOGGER

            LOGGER.setLevel('INFO')
            LOGGER.propagate = 0

            router = Router(self.database, self.app.config['PEEWEE_MIGRATIONS'])
            router.rollback(name)

        @manager.command
        def list():
            """List migrations."""
            from peewee_migrate.router import Router, LOGGER

            LOGGER.setLevel('DEBUG')
            LOGGER.propagate = 0

            router = Router(self.database, self.app.config['PEEWEE_MIGRATIONS'])
            LOGGER.info('Migrations are done:')
            LOGGER.info('\n'.join(router.done))
            LOGGER.info('')
            LOGGER.info('Migrations are undone:')
            LOGGER.info('\n'.join(router.diff))

        return manager
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Flaskify.py
# @author harrydrippin (Seunghwan Hong)
# TODO: .md 파일 생성 기능 추가

import argparse, json, datetime

def flaskify(apifile, comment):
    ret = []
    apifile = json.loads(apifile)

    # Add name, author, doc_comment
    ret.append("# -*- coding: utf-8 -*-")
    ret.append("# app.py")
    ret.append("# Name    : " + apifile['name'])
    ret.append("# Author  : " + apifile['author'])
    ret.append("# Comment : " + apifile['doc_comment'])
    ret.append("\nfrom flask import Flask, request, jsonify")
    ret.append("app = Flask(__name__)")

    ret.append('''
# Exception Region Start

@app.errorhandler(404)
def error_not_found(error):
    return "404"

@app.errorhandler(403)
def error_access_denied(error):
    return "403"

@app.errorhandler(405)
def error_bad_request(error):
    return "405"

@app.errorhandler(500)
def error_internal_error(error):
    return "500"

# Exception Region End
    ''')

    ret.append("# Front API Region Start\n")

    routes = apifile['routes']

    backend_defs = [];
    backend_defs_args = [];
    backend_defs_return = [];

    for route in routes.keys():
        method_save = ""
        if routes[route]['methods'] == None:
            ret.append("@app.route('" + route + "')")
        elif len(routes[route]['methods'].split(",")) >= 1:
            method_temp = "["
            for method in routes[route]['methods'].split(","):
                method_temp += "'" + method + "', "
                method_save += method
            method_temp = method_temp[0:len(method_temp) - 2] + "]"
            ret.append("@app.route('" + route + "', methods=" + method_temp + ")")
        route_code = route.split("/")
        def_name = ""
        for item in route_code:
            def_name += item + "_"
        def_name = def_name[1:len(def_name) - 1]
        backend_defs.append(def_name + "_backend")
        ret.append("def " + def_name + "():")
        if comment == True:
            ret.append("\t# TODO : " + routes[route]['comment'])

        input_dic = routes[route]['input']
        output_dic = routes[route]['output']
        backend_var = ""
        normal_var = ""
        for cursor in input_dic.keys():
            cursor_temp = ""
            normal_var += cursor + ", "
            if method_save.find("POST") != -1:
                cursor_temp = "\t" + cursor + " = request.form['" + cursor + "']"
            elif method_save.find("GET") != -1:
                cursor_temp = "\t" + cursor + " = request.args.get('" + cursor + "', '')"

            if comment == True:
                cursor_temp += " # " + input_dic[cursor]
            ret.append(cursor_temp)

        ret.append("")

        output_var = ""
        output_list = []
        for output in output_dic.keys():
            output_var += output + ", "
            output_list.append(output)

        ret.append("\t" + output_var[0:len(output_var) - 2] + " = " + def_name + "_backend(" + normal_var[0:len(normal_var) - 2] + ")")
        backend_defs_args.append(normal_var[0:len(normal_var) - 2])

        return_text = "\treturn jsonify("

        for txt in output_list:
            return_text += txt + "=" + txt + ", "
        ret.append(return_text[0:len(return_text) - 2] + ")")

        backend_defs_return.append(output_var[0:len(output_var) - 2])

        ret.append("")

    ret.append("# Front API Region End\n\n# Back API Region Start\n")

    for i in range(0, len(backend_defs)):
        ret.append("def " + backend_defs[i] + "(" + backend_defs_args[i] + "):\n\t# TODO : make " + backend_defs[i] + " function\n\treturn " + backend_defs_return[i] + "\n")

    ret.append("# Back API Region End")
    ret.append("\nif __name__ == '__main__':\n\tapp.run(debug=True)")
    return ret

def markdownify(apifile):
    ret = []
    apifile = json.loads(apifile)

    ret.append("# " + apifile['name'])
    ret.append("> 이 파일은 Flaskify에 의하여 자동 생성된 API 문서입니다.")
    ret.append(apifile['doc_comment'])
    ret.append("*****")

    for route in apifile['routes'].keys():
        ret.append("### `" + route + "`")
        ret.append(str(apifile['routes'][route]['comment']))

        input_dic = apifile['routes'][route]['input']
        output_dic = apifile['routes'][route]['output']
        if len(input_dic.keys()) != 0:
            ret.append("#### 입력")
        for inp in input_dic.keys():
            ret.append("* `" + inp + "` — " + input_dic[inp])
        if len(output_dic.keys()) != 0:
            ret.append("#### 출력")
        for outp in output_dic.keys():
            ret.append("* `" + outp + "` — " + output_dic[outp])
        ret.append("*****")

    return ret[0:len(ret) - 1]

def main():
    p = argparse.ArgumentParser(description="ex) flaskify api.txt, flaskify api.txt --comment")
    p.add_argument("apidoc", help="API 문서로의 경로", metavar="[API Document]")
    p.add_argument("-c", "--comment", help="결과에 TODO 주석을 포함합니다.", action="store_true")
    p.add_argument("-m", "--markdown", help="마크다운 문법의 상세 설명 파일을 출력합니다.", action="store_true")

    args = p.parse_args()
    api_doc = args.apidoc
    comment = args.comment
    markdown = args.markdown

    try:
        f = open(api_doc, "rU", encoding="euc-kr")
        apifile = f.readlines()
    except UnicodeDecodeError as e:
        f.close()
        f = open(api_doc, "rU", encoding="utf-8")
        apifile = f.readlines()

    json_string = ""
    for i in apifile:
        json_string += i

    flaskCode = flaskify(json_string, comment)

    try:
        f = open("app.py", "w")
        for i in range(0, len(flaskCode)):
            f.write(flaskCode[i] + "\n")
        f.close()
    except Exception as e:
        print("[-] 예기치 못한 오류가 발생했습니다. 다시 시도하시거나, 에러 내용을 문의해주세요.")
        print("[-] 에러 명세 : " + str(e))
        sys.exit(5)

    if markdown == True:
        markdown_code = markdownify(json_string)
        try:
            f = open("API_DOC.md", "w")
            for i in range(0, len(markdown_code)):
                f.write(markdown_code[i] + "\n")
            f.close()
        except Exception as e:
            print("[-] 예기치 못한 오류가 발생했습니다. 다시 시도하시거나, 에러 내용을 문의해주세요.")
            print("[-] 에러 명세 : " + str(e))
            sys.exit(5)

main()
<commit_msg>Add guide message, shell color<commit_after># -*- coding: utf-8 -*-
# Flaskify.py
# @author harrydrippin (Seunghwan Hong)

import argparse, json, datetime, os

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def flaskify(api_doc, apifile, comment):
    ret = []
    apifile = json.loads(apifile)

    print(bcolors.BOLD + "[+] File Name : " + bcolors.ENDC + api_doc)
    print(bcolors.BOLD + "[+] Name      : " + bcolors.ENDC + apifile['name'])
    print(bcolors.BOLD + "[+] Author    : " + bcolors.ENDC + apifile['author'])
    print(bcolors.BOLD + "[+] Comment   : " + bcolors.ENDC + apifile['doc_comment'] + "\n")
    if comment == True:
        print("[+] Compiling with " + bcolors.HEADER + "TODO " + bcolors.ENDC + "comment...")
    else:
        print("[+] Compiling without " + bcolors.HEADER + "TODO " + bcolors.ENDC + "comment...")

    # Add name, author, doc_comment
    ret.append("# -*- coding: utf-8 -*-")
    ret.append("# app.py")
    ret.append("# Name    : " + apifile['name'])
    ret.append("# Author  : " + apifile['author'])
    ret.append("# Comment : " + apifile['doc_comment'])
    ret.append("\nfrom flask import Flask, request, jsonify")
    ret.append("app = Flask(__name__)")

    ret.append('''
# Exception Region Start

@app.errorhandler(404)
def error_not_found(error):
    return "404"

@app.errorhandler(403)
def error_access_denied(error):
    return "403"

@app.errorhandler(405)
def error_bad_request(error):
    return "405"

@app.errorhandler(500)
def error_internal_error(error):
    return "500"

# Exception Region End
    ''')

    ret.append("# Front API Region Start\n")

    routes = apifile['routes']

    backend_defs = [];
    backend_defs_args = [];
    backend_defs_return = [];

    for route in routes.keys():
        method_save = ""
        if routes[route]['methods'] == None:
            ret.append("@app.route('" + route + "')")
        elif len(routes[route]['methods'].split(",")) >= 1:
            method_temp = "["
            for method in routes[route]['methods'].split(","):
                method_temp += "'" + method + "', "
                method_save += method
            method_temp = method_temp[0:len(method_temp) - 2] + "]"
            ret.append("@app.route('" + route + "', methods=" + method_temp + ")")
        route_code = route.split("/")
        def_name = ""
        for item in route_code:
            def_name += item + "_"
        def_name = def_name[1:len(def_name) - 1]
        backend_defs.append(def_name + "_backend")
        ret.append("def " + def_name + "():")
        if comment == True:
            ret.append("\t# TODO : " + routes[route]['comment'])

        input_dic = routes[route]['input']
        output_dic = routes[route]['output']
        backend_var = ""
        normal_var = ""
        for cursor in input_dic.keys():
            cursor_temp = ""
            normal_var += cursor + ", "
            if method_save.find("POST") != -1:
                cursor_temp = "\t" + cursor + " = request.form['" + cursor + "']"
            elif method_save.find("GET") != -1:
                cursor_temp = "\t" + cursor + " = request.args.get('" + cursor + "', '')"

            if comment == True:
                cursor_temp += " # " + input_dic[cursor]
            ret.append(cursor_temp)

        ret.append("")

        output_var = ""
        output_list = []
        for output in output_dic.keys():
            output_var += output + ", "
            output_list.append(output)

        ret.append("\t" + output_var[0:len(output_var) - 2] + " = " + def_name + "_backend(" + normal_var[0:len(normal_var) - 2] + ")")
        backend_defs_args.append(normal_var[0:len(normal_var) - 2])

        return_text = "\treturn jsonify("

        for txt in output_list:
            return_text += txt + "=" + txt + ", "
        ret.append(return_text[0:len(return_text) - 2] + ")")

        backend_defs_return.append(output_var[0:len(output_var) - 2])

        ret.append("")

    ret.append("# Front API Region End\n\n# Back API Region Start\n")

    for i in range(0, len(backend_defs)):
        ret.append("def " + backend_defs[i] + "(" + backend_defs_args[i] + "):\n\t# TODO : make " + backend_defs[i] + " function\n\treturn " + backend_defs_return[i] + "\n")

    ret.append("# Back API Region End")
    ret.append("\nif __name__ == '__main__':\n\tapp.run(debug=True)")
    return ret

def markdownify(apifile):
    ret = []
    apifile = json.loads(apifile)

    ret.append("# " + apifile['name'])
    ret.append("> 이 파일은 Flaskify에 의하여 자동 생성된 API 문서입니다.")
    ret.append(apifile['doc_comment'])
    ret.append("*****")

    for route in apifile['routes'].keys():
        ret.append("### `" + route + "`")
        ret.append(str(apifile['routes'][route]['comment']))

        input_dic = apifile['routes'][route]['input']
        output_dic = apifile['routes'][route]['output']
        if len(input_dic.keys()) != 0:
            ret.append("#### 입력")
        for inp in input_dic.keys():
            ret.append("* `" + inp + "` — " + input_dic[inp])
        if len(output_dic.keys()) != 0:
            ret.append("#### 출력")
        for outp in output_dic.keys():
            ret.append("* `" + outp + "` — " + output_dic[outp])
        ret.append("*****")

    return ret[0:len(ret) - 1]

def main():
    p = argparse.ArgumentParser(description="ex) flaskify api.txt, flaskify api.txt --comment")
    p.add_argument("apidoc", help="API 문서로의 경로", metavar="[API Document]")
    p.add_argument("-c", "--comment", help="결과에 TODO 주석을 포함합니다.", action="store_true")
    p.add_argument("-m", "--markdown", help="마크다운 문법의 상세 설명 파일을 출력합니다.", action="store_true")

    args = p.parse_args()
    api_doc = args.apidoc
    comment = args.comment
    markdown = args.markdown

    try:
        f = open(api_doc, "rU", encoding="euc-kr")
        apifile = f.readlines()
    except UnicodeDecodeError as e:
        f.close()
        f = open(api_doc, "rU", encoding="utf-8")
        apifile = f.readlines()

    json_string = ""
    for i in apifile:
        json_string += i
    print(bcolors.BOLD + "[*] " + bcolors.BOLD + "Flaskify 1.0.0")
    print("[*] Made by Seunghwan Hong(@harrydrippin on Github)" + bcolors.ENDC + "\n")
    flaskCode = flaskify(api_doc, json_string, comment)

    try:
        f = open("app.py", "w")
        for i in range(0, len(flaskCode)):
            f.write(flaskCode[i] + "\n")
        f.close()
    except Exception as e:
        print("[-] 예기치 못한 오류가 발생했습니다. 다시 시도하시거나, 에러 내용을 문의해주세요.")
        print("[-] 에러 명세 : " + str(e))
        sys.exit(5)

    print("[+] Your " + bcolors.BOLD + "app.py" + bcolors.ENDC + " was saved on " + bcolors.UNDERLINE + os.getcwd() + bcolors.ENDC + ".")

    if markdown == True:
        markdown_code = markdownify(json_string)
        try:
            f = open("API_DOC.md", "w")
            for i in range(0, len(markdown_code)):
                f.write(markdown_code[i] + "\n")
            f.close()
        except Exception as e:
            print("[-] 예기치 못한 오류가 발생했습니다. 다시 시도하시거나, 에러 내용을 문의해주세요.")
            print("[-] 에러 명세 : " + str(e))
            sys.exit(5)
        print("[+] Markdown API document(" + bcolors.BOLD + "API_DOC.md" + bcolors.ENDC + ") was saved. You can also find this on directory above.")

main()
<|endoftext|>


<commit_before>import base64
import logging
from datetime import datetime
from pkg_resources import parse_requirements
from sqlalchemy import UniqueConstraint
from sqlalchemy.ext.associationproxy import association_proxy
from pypi_notifier.extensions import db, github
from pypi_notifier.models.user import User
from pypi_notifier.models.package import Package
from pypi_notifier.models.util import ignored

logger = logging.getLogger(__name__)

class Repo(db.Model):
    __tablename__ = 'repos'
    __table_args__ = (
        UniqueConstraint('user_id', 'github_id'),
    )

    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey(User.id))
    github_id = db.Column(db.Integer)
    name = db.Column(db.String(200))
    last_check = db.Column(db.DateTime)
    last_modified = db.Column(db.String(40))

    packages = association_proxy('requirements', 'package')

    def __init__(self, github_id, user):
        self.github_id = github_id
        self.user = user

    def __repr__(self):
        return "<Repo %s>" % self.name

    @property
    def url(self):
        return "https://github.com/%s" % self.name

    @classmethod
    def update_all_repos(cls):
        repos = cls.query.all()
        for repo in repos:
            with ignored(Exception):
                repo.update_requirements()
                db.session.commit()

    def update_requirements(self):
        """
        Fetches the content of the requirements.txt files from GitHub,
        parses the file and adds each requirement to the repo.

        """
        for project_name, specs in self.parse_requirements_file():
            # specs may be empty list if no version is specified in file
            # No need to add to table since we can't check updates.
            if specs:
                # There must be '==' operator in specs.
                operators = [s[0] for s in specs]
                if '==' in operators:
                    # If the project is not registered on PyPI,
                    # we are not adding it.
                    if project_name.lower() in Package.get_all_names():
                        self.add_new_requirement(project_name, specs)

        self.last_check = datetime.utcnow()

    def add_new_requirement(self, name, specs):
        from pypi_notifier.models.requirement import Requirement
        package = db.session.query(Package).filter(Package.name == name).first()
        if not package:
            package = Package(name=name)
            db.session.add(package)
            db.session.flush([package])

        requirement = db.session.query(Requirement).filter(
                Requirement.repo_id == self.id,
                Requirement.package_id == package.id).first()
        if not requirement:
            requirement = Requirement(repo=self, package=package)

        requirement.specs = specs
        self.requirements.append = requirement

    def parse_requirements_file(self):
        contents = self.fetch_requirements()
        if contents:
            contents = strip_requirements(contents)
            if contents:
                for req in parse_requirements(contents):
                    yield req.project_name.lower(), req.specs

    def fetch_requirements(self):
        logger.info("Fetching requirements of repo: %s", self)
        path = 'repos/%s/contents/requirements.txt' % self.name
        headers = None
        if self.last_modified:
            headers = {'If-Modified-Since': self.last_modified}

        response = github.raw_request('GET', path, headers=headers, access_token=self.user.github_token)
        logger.debug("Response: %s", response)
        if response.status_code == 200:
            self.last_modified = response.headers['Last-Modified']
            response = response.json()
            if response['encoding'] == 'base64':
                return base64.b64decode(response['content']).decode('utf-8', 'replace')
            else:
                raise Exception("Unknown encoding: %s" % response['encoding'])
        elif response.status_code == 304:  # Not modified
            return None
        elif response.status_code == 401:
            # User's token is not valid. Let's delete the user.
            db.session.delete(self.user)
        elif response.status_code == 404:
            # requirements.txt file is not found.
            # Remove the repo so we won't check it again.
            db.session.delete(self)
        else:
            raise Exception("Unknown status code: %s" % response.status_code)

def strip_requirements(s):
    """
    Cleans up requirements.txt contents and returns as new str.

    pkg_resources.parse_requirements() cannot parse the file if it contains
    an option for index URL.
        Example: "-i http://simple.crate.io/"

    Also it cannot parse the repository URLs.
        Example: git+https://github.com/pythonforfacebook/facebook-sdk.git

    """
    ignore_lines = (
        '-e',  # editable
        '-i', '--index-url',  # other source
        'git+', 'svn+', 'hg+', 'bzr+',  # vcs
        '-r',  # include other files (not supported yet) TODO
    )
    return '\n'.join(l for l in s.splitlines()
                     if not l.strip().startswith(ignore_lines))
<commit_msg>refactor<commit_after>import base64
import logging
from datetime import datetime
from pkg_resources import parse_requirements
from sqlalchemy import UniqueConstraint
from sqlalchemy.ext.associationproxy import association_proxy
from pypi_notifier.extensions import db, github
from pypi_notifier.models.user import User
from pypi_notifier.models.package import Package
from pypi_notifier.models.util import ignored

logger = logging.getLogger(__name__)

class Repo(db.Model):
    __tablename__ = 'repos'
    __table_args__ = (
        UniqueConstraint('user_id', 'github_id'),
    )

    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey(User.id))
    github_id = db.Column(db.Integer)
    name = db.Column(db.String(200))
    last_check = db.Column(db.DateTime)
    last_modified = db.Column(db.String(40))

    packages = association_proxy('requirements', 'package')

    def __init__(self, github_id, user):
        self.github_id = github_id
        self.user = user

    def __repr__(self):
        return "<Repo %s>" % self.name

    @property
    def url(self):
        return "https://github.com/%s" % self.name

    @classmethod
    def update_all_repos(cls):
        repos = cls.query.all()
        for repo in repos:
            with ignored(Exception):
                repo.update_requirements()
                db.session.commit()

    def update_requirements(self):
        """
        Fetches the content of the requirements.txt files from GitHub,
        parses the file and adds each requirement to the repo.

        """
        for project_name, specs in self.parse_requirements_file():
            # specs may be empty list if no version is specified in file
            # No need to add to table since we can't check updates.
            if specs:
                # There must be '==' operator in specs.
                operators = [s[0] for s in specs]
                if '==' in operators:
                    # If the project is not registered on PyPI,
                    # we are not adding it.
                    if project_name.lower() in Package.get_all_names():
                        self.add_new_requirement(project_name, specs)

        self.last_check = datetime.utcnow()

    def add_new_requirement(self, name, specs):
        from pypi_notifier.models.requirement import Requirement
        package = db.session.query(Package).filter(Package.name == name).first()
        if not package:
            package = Package(name=name)
            db.session.add(package)
            db.session.flush([package])

        requirement = db.session.query(Requirement).filter(
                Requirement.repo_id == self.id,
                Requirement.package_id == package.id).first()
        if not requirement:
            requirement = Requirement(repo=self, package=package)

        requirement.specs = specs
        self.requirements.append = requirement

    def parse_requirements_file(self):
        contents = self.fetch_requirements()
        if contents:
            contents = strip_requirements(contents)
            if contents:
                for req in parse_requirements(contents):
                    yield req.project_name.lower(), req.specs

    def fetch_requirements(self):
        logger.info("Fetching requirements of repo: %s", self)
        path = 'repos/%s/contents/requirements.txt' % self.name
        headers = None
        if self.last_modified:
            headers = {'If-Modified-Since': self.last_modified}

        response = github.raw_request('GET', path, headers=headers, access_token=self.user.github_token)
        logger.debug("Response: %s", response)
        if response.status_code == 200:
            self.last_modified = response.headers['Last-Modified']
            response = response.json()
            if response['encoding'] == 'base64':
                return base64.b64decode(response['content']).decode('utf-8', 'replace')
            else:
                raise Exception("Unknown encoding: %s" % response['encoding'])
        elif response.status_code == 304:  # Not modified
            return None
        elif response.status_code == 401:
            # User's token is not valid. Let's delete the user.
            db.session.delete(self.user)
        elif response.status_code == 404:
            # requirements.txt file is not found.
            # Remove the repo so we won't check it again.
            db.session.delete(self)
        else:
            raise Exception("Unknown status code: %s" % response.status_code)

def strip_requirements(s):
    """
    Cleans up requirements.txt contents and returns as new str.

    pkg_resources.parse_requirements() cannot parse the file if it contains
    an option for index URL.
        Example: "-i http://simple.crate.io/"

    Also it cannot parse the repository URLs.
        Example: git+https://github.com/pythonforfacebook/facebook-sdk.git

    """
    ignore_lines = (
        '-e',  # editable
        '-i', '--index-url',  # other source
        'git+', 'svn+', 'hg+', 'bzr+',  # vcs
        '-r',  # include other files (not supported yet) TODO
    )
    return '\n'.join(l for l in s.splitlines() if not l.strip().startswith(ignore_lines))
<|endoftext|>


<commit_before>from pyrate.repositories import sql
import psycopg2
import logging
try:
    import pandas as pd
except ImportError:
    logging.warn("No pandas found")
    pd = None

EXPORT_COMMANDS = [('status', 'report status of this repository.'),
                   ('create', 'create the repository.'),
                   ('truncate', 'delete all data in this repository.')]

def load(options, readonly=False):
    return AISdb(options, readonly)

class AISdb(sql.PgsqlRepository):

    double_type = 'double precision'

    clean_db_spec = {
        'cols': [
            ('MMSI', 'integer'),
            ('Time', 'timestamp without time zone'),
            ('Message_ID', 'integer'),
            ('Navigational_status', 'integer'),
            ('SOG', double_type),
            ('Longitude', double_type),
            ('Latitude', double_type),
            ('COG', double_type),
            ('Heading', double_type),
            ('IMO', 'integer null'),
            ('Draught', double_type),
            ('Destination', 'character varying(255)'),
            ('Vessel_Name', 'character varying(255)'),
            ('ETA_month', 'integer'),
            ('ETA_day', 'integer'),
            ('ETA_hour', 'integer'),
            ('ETA_minute', 'integer'),
            ('source', 'smallint'),
            ('ID', 'SERIAL PRIMARY KEY')
        ],
        'indices': [
            ('dt_idx', ['Time']),
            ('imo_idx', ['IMO']),
            ('lonlat_idx', ['Longitude', 'Latitude']),
            ('mmsi_idx', ['MMSI']),
            ('msg_idx', ['Message_ID']),
            ('source_idx', ['source']),
            ('mmsi_imo_idx', ['MMSI','IMO'])
        ]
    }

    dirty_db_spec = {
        'cols': [
            ('MMSI', 'bigint'),
            ('Time', 'timestamp without time zone'),
            ('Message_ID', 'integer'),
            ('Navigational_status', 'integer'),
            ('SOG', double_type),
            ('Longitude', double_type),
            ('Latitude', double_type),
            ('COG', double_type),
            ('Heading', double_type),
            ('IMO', 'integer null'),
            ('Draught', double_type),
            ('Destination', 'character varying(255)'),
            ('Vessel_Name', 'character varying(255)'),
            ('ETA_month', 'integer'),
            ('ETA_day', 'integer'),
            ('ETA_hour', 'integer'),
            ('ETA_minute', 'integer'),
            ('source', 'smallint'),
            ('ID', 'SERIAL PRIMARY KEY')
        ],
        'indices': [
            ('dt_idx', ['Time']),
            ('imo_idx', ['IMO']),
            ('lonlat_idx', ['Longitude', 'Latitude']),
            ('mmsi_idx', ['MMSI']),
            ('msg_idx', ['Message_ID']),
            ('source_idx', ['source']),
            ('mmsi_imo_idx', ['MMSI','IMO'])
        ]
    }

    sources_db_spec = {
        'cols': [
            ('ID', 'SERIAL PRIMARY KEY'),
            ('timestamp', 'timestamp without time zone DEFAULT now()'),
            ('filename', 'TEXT'),
            ('ext', 'TEXT'),
            ('invalid', 'integer'),
            ('clean', 'integer'),
            ('dirty', 'integer'),
            ('source', 'integer')
        ]
    }

    imolist_db_spec = {
        'cols': [
            ('mmsi', 'integer NOT NULL'),
            ('imo', 'integer NULL'),
            ('first_seen', 'timestamp without time zone'),
            ('last_seen', 'timestamp without time zone')
        ],
        'constraint': ['CONSTRAINT imo_list_key UNIQUE (mmsi, imo)']
    }

    clean_imo_list = {
        'cols': imolist_db_spec['cols'],
        'constraint': ['CONSTRAINT imo_list_pkey PRIMARY KEY (mmsi, imo)']
    }

    action_log_spec = {
        'cols': [
            ('timestamp', 'timestamp without time zone DEFAULT now()'),
            ('action', 'TEXT'),
            ('mmsi', 'integer NOT NULL'),
            ('ts_from', 'timestamp without time zone'),
            ('ts_to', 'timestamp without time zone'),
            ('count', 'integer NULL')
        ],
        'indices': [
            ('ts_idx', ['timestamp']),
            ('action_idx', ['action']),
            ('mmsi_idx', ['mmsi'])
        ],
        'constraint': ['CONSTRAINT action_log_pkey PRIMARY KEY (timestamp, action, mmsi)']
    }

    def __init__(self, options, readonly=False):
        super(AISdb, self).__init__(options, readonly)
        self.clean = sql.Table(self, 'ais_clean', self.clean_db_spec['cols'],
                               self.clean_db_spec['indices'])
        self.dirty = sql.Table(self, 'ais_dirty', self.dirty_db_spec['cols'],
                               self.dirty_db_spec['indices'])
        self.sources = sql.Table(self, 'ais_sources', self.sources_db_spec['cols'])
        self.imolist = sql.Table(self, 'imo_list', self.imolist_db_spec['cols'],
                                 constraint=self.imolist_db_spec['constraint'])
        self.extended = AISExtendedTable(self)
        self.clean_imolist = sql.Table(self, 'imo_list_clean', self.clean_imo_list['cols'], constraint=self.clean_imo_list['constraint'])
        self.action_log = sql.Table(self, 'action_log', self.action_log_spec['cols'], self.action_log_spec['indices'], constraint=self.action_log_spec['constraint'])
        self.tables = [self.clean, self.dirty, self.sources, self.imolist, self.extended, self.clean_imolist, self.action_log]

    def status(self):
        print("Status of PGSql database "+ self.db +":")
        for tb in self.tables:
            s = tb.status()
            if s >= 0:
                print("Table {}: {} rows.".format(tb.get_name(), s))
            else:
                print("Table {}: not yet created.".format(tb.get_name()))

    def create(self):
        """Create the tables for the AIS data."""
        for tb in self.tables:
            tb.create()

    def truncate(self):
        """Delete all data in the AIS table."""
        for tb in self.tables:
            tb.truncate()

    def ship_info(self, imo):
        with self.conn.cursor() as cur:
            cur.execute("select vessel_name, MIN(time), MAX(time) from ais_clean where message_id = 5 and imo = %s GROUP BY vessel_name", [imo])
            for row in cur:
                print("Vessel: {} ({} - {})".format(*row))

            cur.execute("select mmsi, first_seen, last_seen from imo_list where imo = %s", [imo])
            for row in cur:
                print("MMSI = {} ({} - {})".format(*row))

    def get_messages_for_vessel(self, imo, from_ts=None, to_ts=None, use_clean_db=False, as_df=False):
        if use_clean_db:
            imo_list = self.imolist
        else:
            imo_list = self.clean_imolist

        where = ["imo = {}"]
        params = [imo]
        #Amended EOK - no time field in this table
        # if not from_ts is None:
        #     where.append("time >= {}")
        #     params.append(from_ts)
        # if not to_ts is None:
        #     where.append("time <= {}")
        #     params.append(to_ts)

        with self.conn.cursor() as cur:
            cur.execute("select mmsi, first_seen, last_seen from {} where {}".format(imo_list.name, ' AND '.join(where)).format(*params))
            msg_stream = None
            # get data for each of this ship's mmsi numbers, and concat
            for mmsi, first, last in cur:
                stream = self.get_message_stream(mmsi, from_ts=first, to_ts=last, use_clean_db=use_clean_db, as_df=as_df)
                if msg_stream is None:
                    msg_stream = stream
                else:
                    msg_stream = msg_stream + stream
            return msg_stream

    def get_message_stream(self, mmsi, from_ts=None, to_ts=None, use_clean_db=False, as_df=False):
        """Gets the stream of messages for the given mmsi, ordered by timestamp ascending"""
        # construct db query
        if use_clean_db:
            db = self.clean
        else:
            db = self.extended
        where = ["mmsi = %s"]
        params = [mmsi]
        if not from_ts is None:
            where.append("time >= %s")
            params.append(from_ts)
        if not to_ts is None:
            where.append("time <= %s")
            params.append(to_ts)

        cols_list = ','.join([c[0].lower() for c in db.cols])
        where_clause = ' AND '.join(where)
        sql = "SELECT {} FROM {} WHERE {} ORDER BY time ASC".format(cols_list,
                db.get_name(), where_clause)

        if as_df:
            if pd is None:
                raise RuntimeError("Pandas not found, cannot create dataframe")
            # create pandas dataframe
            with self.conn.cursor() as cur:
                full_sql = cur.mogrify(sql, params).decode('ascii')
            return pd.read_sql(full_sql, self.conn, index_col='time', parse_dates=['time'])

        else:
            with self.conn.cursor() as cur:
                cur.execute(sql, params)
                msg_stream = []
                # convert tuples from db cursor into dicts
                for row in cur:
                    message = {}
                    for i, col in enumerate(db.cols):
                        message[col[0]] = row[i]
                    msg_stream.append(message)

                return msg_stream

class AISExtendedTable(sql.Table):

    def __init__(self, db):
        super(AISExtendedTable, self).__init__(db, 'ais_extended',
            AISdb.clean_db_spec['cols'] + [('location', 'geography(POINT, 4326)')],
            AISdb.clean_db_spec['indices'])

    def create(self):
        with self.db.conn.cursor() as cur:
            cur.execute("CREATE EXTENSION IF NOT EXISTS postgis")
        super(AISExtendedTable, self).create()
        with self.db.conn.cursor() as cur:
            # trigger for GIS location generation
            cur.execute("""CREATE OR REPLACE FUNCTION location_insert() RETURNS trigger AS '
                        BEGIN
                            NEW."location" := ST_SetSRID(ST_MakePoint(NEW.longitude, NEW.latitude),4326);
                            RETURN NEW;
                        END;
                        ' LANGUAGE plpgsql;
                        CREATE TRIGGER {0}_gis_insert
                        BEFORE INSERT OR UPDATE ON {0} FOR EACH ROW EXECUTE PROCEDURE location_insert();
                        """.format(self.name))
        self.db.conn.commit()

    def create_indices(self):
        with self.db.conn.cursor() as cur:
            idxn = self.name.lower() + "_location_idx"
            try:
                logging.info("CREATING GIST INDEX "+ idxn + " on table "+ self.name)
                cur.execute("CREATE INDEX \""+ idxn +"\" ON \"" + self.name +"\" USING GIST(\"location\")")
            except psycopg2.ProgrammingError:
                logging.info("Index "+ idxn +" already exists")
                self.db.conn.rollback()
        super(AISExtendedTable, self).create_indices()

    def drop_indices(self):
        with self.db.conn.cursor() as cur:
            tbl = self.name
            idxn = tbl.lower() + "_location_idx"
            logging.info("Dropping index: "+ idxn + " on table "+ tbl)
            cur.execute("DROP INDEX IF EXISTS \""+ idxn +"\"")
        super(AISExtendedTable, self).drop_indices()
<commit_msg>If a table already existed, the the trigger generation failed and caused it all to fall over. This isn't a problem if you're creating from scratch but it is if you're looking to update an existing database by adding new tables. I've used the same aprroach as has been done already in pyrate, by surrounding in try/catch.<commit_after>from pyrate.repositories import sql
import psycopg2
import logging
try:
    import pandas as pd
except ImportError:
    logging.warn("No pandas found")
    pd = None

EXPORT_COMMANDS = [('status', 'report status of this repository.'),
                   ('create', 'create the repository.'),
                   ('truncate', 'delete all data in this repository.')]

def load(options, readonly=False):
    return AISdb(options, readonly)

class AISdb(sql.PgsqlRepository):

    double_type = 'double precision'

    clean_db_spec = {
        'cols': [
            ('MMSI', 'integer'),
            ('Time', 'timestamp without time zone'),
            ('Message_ID', 'integer'),
            ('Navigational_status', 'integer'),
            ('SOG', double_type),
            ('Longitude', double_type),
            ('Latitude', double_type),
            ('COG', double_type),
            ('Heading', double_type),
            ('IMO', 'integer null'),
            ('Draught', double_type),
            ('Destination', 'character varying(255)'),
            ('Vessel_Name', 'character varying(255)'),
            ('ETA_month', 'integer'),
            ('ETA_day', 'integer'),
            ('ETA_hour', 'integer'),
            ('ETA_minute', 'integer'),
            ('source', 'smallint'),
            ('ID', 'SERIAL PRIMARY KEY')
        ],
        'indices': [
            ('dt_idx', ['Time']),
            ('imo_idx', ['IMO']),
            ('lonlat_idx', ['Longitude', 'Latitude']),
            ('mmsi_idx', ['MMSI']),
            ('msg_idx', ['Message_ID']),
            ('source_idx', ['source']),
            ('mmsi_imo_idx', ['MMSI','IMO'])
        ]
    }

    dirty_db_spec = {
        'cols': [
            ('MMSI', 'bigint'),
            ('Time', 'timestamp without time zone'),
            ('Message_ID', 'integer'),
            ('Navigational_status', 'integer'),
            ('SOG', double_type),
            ('Longitude', double_type),
            ('Latitude', double_type),
            ('COG', double_type),
            ('Heading', double_type),
            ('IMO', 'integer null'),
            ('Draught', double_type),
            ('Destination', 'character varying(255)'),
            ('Vessel_Name', 'character varying(255)'),
            ('ETA_month', 'integer'),
            ('ETA_day', 'integer'),
            ('ETA_hour', 'integer'),
            ('ETA_minute', 'integer'),
            ('source', 'smallint'),
            ('ID', 'SERIAL PRIMARY KEY')
        ],
        'indices': [
            ('dt_idx', ['Time']),
            ('imo_idx', ['IMO']),
            ('lonlat_idx', ['Longitude', 'Latitude']),
            ('mmsi_idx', ['MMSI']),
            ('msg_idx', ['Message_ID']),
            ('source_idx', ['source']),
            ('mmsi_imo_idx', ['MMSI','IMO'])
        ]
    }

    sources_db_spec = {
        'cols': [
            ('ID', 'SERIAL PRIMARY KEY'),
            ('timestamp', 'timestamp without time zone DEFAULT now()'),
            ('filename', 'TEXT'),
            ('ext', 'TEXT'),
            ('invalid', 'integer'),
            ('clean', 'integer'),
            ('dirty', 'integer'),
            ('source', 'integer')
        ]
    }

    imolist_db_spec = {
        'cols': [
            ('mmsi', 'integer NOT NULL'),
            ('imo', 'integer NULL'),
            ('first_seen', 'timestamp without time zone'),
            ('last_seen', 'timestamp without time zone')
        ],
        'constraint': ['CONSTRAINT imo_list_key UNIQUE (mmsi, imo)']
    }

    clean_imo_list = {
        'cols': imolist_db_spec['cols'],
        'constraint': ['CONSTRAINT imo_list_pkey PRIMARY KEY (mmsi, imo)']
    }

    action_log_spec = {
        'cols': [
            ('timestamp', 'timestamp without time zone DEFAULT now()'),
            ('action', 'TEXT'),
            ('mmsi', 'integer NOT NULL'),
            ('ts_from', 'timestamp without time zone'),
            ('ts_to', 'timestamp without time zone'),
            ('count', 'integer NULL')
        ],
        'indices': [
            ('ts_idx', ['timestamp']),
            ('action_idx', ['action']),
            ('mmsi_idx', ['mmsi'])
        ],
        'constraint': ['CONSTRAINT action_log_pkey PRIMARY KEY (timestamp, action, mmsi)']
    }

    def __init__(self, options, readonly=False):
        super(AISdb, self).__init__(options, readonly)
        self.clean = sql.Table(self, 'ais_clean', self.clean_db_spec['cols'],
                               self.clean_db_spec['indices'])
        self.dirty = sql.Table(self, 'ais_dirty', self.dirty_db_spec['cols'],
                               self.dirty_db_spec['indices'])
        self.sources = sql.Table(self, 'ais_sources', self.sources_db_spec['cols'])
        self.imolist = sql.Table(self, 'imo_list', self.imolist_db_spec['cols'],
                                 constraint=self.imolist_db_spec['constraint'])
        self.extended = AISExtendedTable(self)
        self.clean_imolist = sql.Table(self, 'imo_list_clean', self.clean_imo_list['cols'], constraint=self.clean_imo_list['constraint'])
        self.action_log = sql.Table(self, 'action_log', self.action_log_spec['cols'], self.action_log_spec['indices'], constraint=self.action_log_spec['constraint'])
        self.tables = [self.clean, self.dirty, self.sources, self.imolist, self.extended, self.clean_imolist, self.action_log]

    def status(self):
        print("Status of PGSql database "+ self.db +":")
        for tb in self.tables:
            s = tb.status()
            if s >= 0:
                print("Table {}: {} rows.".format(tb.get_name(), s))
            else:
                print("Table {}: not yet created.".format(tb.get_name()))

    def create(self):
        """Create the tables for the AIS data."""
        for tb in self.tables:
            tb.create()

    def truncate(self):
        """Delete all data in the AIS table."""
        for tb in self.tables:
            tb.truncate()

    def ship_info(self, imo):
        with self.conn.cursor() as cur:
            cur.execute("select vessel_name, MIN(time), MAX(time) from ais_clean where message_id = 5 and imo = %s GROUP BY vessel_name", [imo])
            for row in cur:
                print("Vessel: {} ({} - {})".format(*row))

            cur.execute("select mmsi, first_seen, last_seen from imo_list where imo = %s", [imo])
            for row in cur:
                print("MMSI = {} ({} - {})".format(*row))

    def get_messages_for_vessel(self, imo, from_ts=None, to_ts=None, use_clean_db=False, as_df=False):
        if use_clean_db:
            imo_list = self.imolist
        else:
            imo_list = self.clean_imolist

        where = ["imo = {}"]
        params = [imo]
        #Amended EOK - no time field in this table
        # if not from_ts is None:
        #     where.append("time >= {}")
        #     params.append(from_ts)
        # if not to_ts is None:
        #     where.append("time <= {}")
        #     params.append(to_ts)

        with self.conn.cursor() as cur:
            cur.execute("select mmsi, first_seen, last_seen from {} where {}".format(imo_list.name, ' AND '.join(where)).format(*params))
            msg_stream = None
            # get data for each of this ship's mmsi numbers, and concat
            for mmsi, first, last in cur:
                stream = self.get_message_stream(mmsi, from_ts=first, to_ts=last, use_clean_db=use_clean_db, as_df=as_df)
                if msg_stream is None:
                    msg_stream = stream
                else:
                    msg_stream = msg_stream + stream
            return msg_stream

    def get_message_stream(self, mmsi, from_ts=None, to_ts=None, use_clean_db=False, as_df=False):
        """Gets the stream of messages for the given mmsi, ordered by timestamp ascending"""
        # construct db query
        if use_clean_db:
            db = self.clean
        else:
            db = self.extended
        where = ["mmsi = %s"]
        params = [mmsi]
        if not from_ts is None:
            where.append("time >= %s")
            params.append(from_ts)
        if not to_ts is None:
            where.append("time <= %s")
            params.append(to_ts)

        cols_list = ','.join([c[0].lower() for c in db.cols])
        where_clause = ' AND '.join(where)
        sql = "SELECT {} FROM {} WHERE {} ORDER BY time ASC".format(cols_list,
                db.get_name(), where_clause)

        if as_df:
            if pd is None:
                raise RuntimeError("Pandas not found, cannot create dataframe")
            # create pandas dataframe
            with self.conn.cursor() as cur:
                full_sql = cur.mogrify(sql, params).decode('ascii')
            return pd.read_sql(full_sql, self.conn, index_col='time', parse_dates=['time'])

        else:
            with self.conn.cursor() as cur:
                cur.execute(sql, params)
                msg_stream = []
                # convert tuples from db cursor into dicts
                for row in cur:
                    message = {}
                    for i, col in enumerate(db.cols):
                        message[col[0]] = row[i]
                    msg_stream.append(message)

                return msg_stream

class AISExtendedTable(sql.Table):

    def __init__(self, db):
        super(AISExtendedTable, self).__init__(db, 'ais_extended',
            AISdb.clean_db_spec['cols'] + [('location', 'geography(POINT, 4326)')],
            AISdb.clean_db_spec['indices'])

    def create(self):
        with self.db.conn.cursor() as cur:
            cur.execute("CREATE EXTENSION IF NOT EXISTS postgis")
        super(AISExtendedTable, self).create()
        with self.db.conn.cursor() as cur:
            # trigger for GIS location generation
            try:
                cur.execute("""CREATE OR REPLACE FUNCTION location_insert() RETURNS trigger AS '
                        BEGIN
                            NEW."location" := ST_SetSRID(ST_MakePoint(NEW.longitude, NEW.latitude),4326);
                            RETURN NEW;
                        END;
                        ' LANGUAGE plpgsql;
                        CREATE TRIGGER {0}_gis_insert
                        BEFORE INSERT OR UPDATE ON {0} FOR EACH ROW EXECUTE PROCEDURE location_insert();
                        """.format(self.name))
            except psycopg2.ProgrammingError:
                logging.info("{}_gis_insert already exists".format(self.name))
                self.db.conn.rollback()
        self.db.conn.commit()

    def create_indices(self):
        with self.db.conn.cursor() as cur:
            idxn = self.name.lower() + "_location_idx"
            try:
                logging.info("CREATING GIST INDEX "+ idxn + " on table "+ self.name)
                cur.execute("CREATE INDEX \""+ idxn +"\" ON \"" + self.name +"\" USING GIST(\"location\")")
            except psycopg2.ProgrammingError:
                logging.info("Index "+ idxn +" already exists")
                self.db.conn.rollback()
        super(AISExtendedTable, self).create_indices()

    def drop_indices(self):
        with self.db.conn.cursor() as cur:
            tbl = self.name
            idxn = tbl.lower() + "_location_idx"
            logging.info("Dropping index: "+ idxn + " on table "+ tbl)
            cur.execute("DROP INDEX IF EXISTS \""+ idxn +"\"")
        super(AISExtendedTable, self).drop_indices()
<|endoftext|>


<commit_before>import io
import os
import struct
import socket
import threading

from pyroute2.iocore.addrpool import AddrPool  # FIXME: move to common
from pyroute2.netlink import nlmsg
from pyroute2.netlink import mtypes
from pyroute2.netlink import NetlinkError
from pyroute2.netlink import NetlinkDecodeError
from pyroute2.netlink import NetlinkHeaderDecodeError
from pyroute2.netlink import NLMSG_ERROR
from pyroute2.netlink import NETLINK_GENERIC

class Marshal(object):
    '''
    Generic marshalling class
    '''

    msg_map = {}
    debug = False

    def __init__(self):
        self.lock = threading.Lock()
        # one marshal instance can be used to parse one
        # message at once
        self.msg_map = self.msg_map or {}
        self.defragmentation = {}

    def parse(self, data, sock=None):
        '''
        Parse the data in the buffer

        If socket is provided, support defragmentation
        '''
        with self.lock:
            data.seek(0)
            offset = 0
            result = []

            if sock in self.defragmentation:
                save = self.defragmentation[sock]
                save.write(data.read())
                save.length += data.length
                # discard save
                data = save
                del self.defragmentation[sock]
                data.seek(0)

            while offset < data.length:
                # pick type and length
                (length, msg_type) = struct.unpack('IH', data.read(6))
                data.seek(offset)
                # if length + offset is greater than
                # remaining size, save the buffer for
                # defragmentation
                if (sock is not None) and (length + offset > data.length):
                    # create save buffer
                    self.defragmentation[sock] = save = io.BytesIO()
                    save.length = save.write(data.read())
                    # truncate data
                    data.truncate(offset)
                    break

                error = None
                if msg_type == NLMSG_ERROR:
                    data.seek(offset + 16)
                    code = abs(struct.unpack('i', data.read(4))[0])
                    if code > 0:
                        error = NetlinkError(code)
                    data.seek(offset)

                msg_class = self.msg_map.get(msg_type, nlmsg)
                msg = msg_class(data, debug=self.debug)
                try:
                    msg.decode()
                    msg['header']['error'] = error
                    # try to decode encapsulated error message
                    if error is not None:
                        enc_type = struct.unpack('H', msg.raw[24:26])[0]
                        enc_class = self.msg_map.get(enc_type, nlmsg)
                        enc = enc_class(msg.raw[20:])
                        enc.decode()
                        msg['header']['errmsg'] = enc
                except NetlinkHeaderDecodeError as e:
                    # in the case of header decoding error,
                    # create an empty message
                    msg = nlmsg()
                    msg['header']['error'] = e
                except NetlinkDecodeError as e:
                    msg['header']['error'] = e
                mtype = msg['header'].get('type', None)
                if mtype in (1, 2, 3, 4):
                    msg['event'] = mtypes.get(mtype, 'none')
                self.fix_message(msg)
                offset += msg.length
                result.append(msg)

            return result

    def fix_message(self, msg):
        pass

# 8<-----------------------------------------------------------
# Singleton, containing possible modifiers to the NetlinkSocket
# bind() call.
#
# Normally, you can open only one netlink connection for one
# process, but there is a hack. Current PID_MAX_LIMIT is 2^22,
# so we can use the rest to midify pid field.
#
# See also libnl library, lib/socket.c:generate_local_port()
sockets = AddrPool(minaddr=0x0,
                   maxaddr=0x3ff,
                   reverse=True)
# 8<-----------------------------------------------------------

class NetlinkSocket(socket.socket):
    '''
    Generic netlink socket
    '''

    def __init__(self, family=NETLINK_GENERIC, port=None, pid=None):
        socket.socket.__init__(self, socket.AF_NETLINK,
                               socket.SOCK_DGRAM, family)
        global sockets

        # 8<-----------------------------------------
        # PID init is here only for compatibility,
        # later it will be completely moved to bind()
        self.epid = None
        self.port = 0
        self.fixed = True
        if pid is None:
            self.pid = os.getpid() & 0x3fffff
            self.port = port
            self.fixed = self.port is not None
        elif pid == 0:
            self.pid = os.getpid()
        else:
            self.pid = pid
        # 8<-----------------------------------------
        self.groups = 0
        self.marshal = Marshal()

    def register_policy(self, policy, msg_class=None):
        '''
        Register netlink encoding/decoding policy. Can
        be specified in two ways:
        `nlsocket.register_policy(MSG_ID, msg_class)`
        to register one particular rule, or
        `nlsocket.register_policy({MSG_ID1: msg_class})`
        to register several rules at once.
        E.g.::

            policy = {RTM_NEWLINK: ifinfmsg,
                      RTM_DELLINK: ifinfmsg,
                      RTM_NEWADDR: ifaddrmsg,
                      RTM_DELADDR: ifaddrmsg}
            nlsocket.register_policy(policy)

        One can call `register_policy()` as many times,
        as one want to -- it will just extend the current
        policy scheme, not replace it.
        '''
        if isinstance(policy, int) and msg_class is not None:
            policy = {policy: msg_class}

        assert isinstance(policy, dict)
        for key in policy:
            self.marshal.msg_map[key] = policy[key]

        return self.marshal.msg_map

    def unregister_policy(self, policy):
        '''
        Unregister policy. Policy can be:

        * int -- then it will just remove one policy
        * list or tuple of ints -- remove all given
        * dict -- remove policies by keys from dict

        In the last case the routine will ignore dict values,
        it is implemented so just to make it compatible with
        `get_policy_map()` return value.
        '''
        if isinstance(policy, int):
            policy = [policy]
        elif isinstance(policy, dict):
            policy = list(policy)

        assert isinstance(policy, (tuple, list, set))

        for key in policy:
            del self.marshal.msg_map[key]

        return self.marshal.msg_map

    def get_policy_map(self, policy=None):
        '''
        Return policy for a given message type or for all
        message types. Policy parameter can be either int,
        or a list of ints. Always return dictionary.
        '''
        if policy is None:
            return self.marshal.msg_map

        if isinstance(policy, int):
            policy = [policy]

        assert isinstance(policy, (list, tuple, set))

        ret = {}
        for key in policy:
            ret[key] = self.marshal.msg_map[key]

        return ret

    def bind(self, groups=0, pid=None):
        '''
        Bind the socket to given multicast groups, using
        given pid.

        * If pid is None, use automatic port allocation
        * If pid == 0, use process' pid
        * If pid == <int>, use the value instead of pid
        '''
        if pid is not None:
            self.port = 0
            self.fixed = True
            self.pid = pid or os.getpid()

        self.groups = groups
        # if we have pre-defined port, use it strictly
        if self.fixed:
            self.epid = self.pid + (self.port << 22)
            socket.socket.bind(self, (self.epid, self.groups))
            return

        # if we have no pre-defined port, scan all the
        # range till the first available port
        for i in range(1024):
            try:
                self.port = sockets.alloc()
                self.epid = self.pid + (self.port << 22)
                socket.socket.bind(self, (self.epid, self.groups))
                # if we're here, bind() done successfully, just exit
                return
            except socket.error as e:
                # pass occupied sockets, raise other exceptions
                if e.errno != 98:
                    raise
        else:
            # raise "address in use" -- to be compatible
            raise socket.error(98, 'Address already in use')

    def get(self):
        '''
        Get parsed messages list.
        '''
        data = io.BytesIO()
        data.length = data.write(self.recv(16384))
        return self.marshal.parse(data, self)

    def close(self):
        '''
        Correctly close the socket and free all resources.
        '''
        global sockets
        if self.epid is not None:
            assert self.port is not None
            if not self.fixed:
                sockets.free(self.port)
            self.epid = None
        socket.socket.close(self)
<commit_msg>nlsocket: expose bufsize in `get()` call<commit_after>import io
import os
import struct
import socket
import threading

from pyroute2.iocore.addrpool import AddrPool  # FIXME: move to common
from pyroute2.netlink import nlmsg
from pyroute2.netlink import mtypes
from pyroute2.netlink import NetlinkError
from pyroute2.netlink import NetlinkDecodeError
from pyroute2.netlink import NetlinkHeaderDecodeError
from pyroute2.netlink import NLMSG_ERROR
from pyroute2.netlink import NETLINK_GENERIC

class Marshal(object):
    '''
    Generic marshalling class
    '''

    msg_map = {}
    debug = False

    def __init__(self):
        self.lock = threading.Lock()
        # one marshal instance can be used to parse one
        # message at once
        self.msg_map = self.msg_map or {}
        self.defragmentation = {}

    def parse(self, data, sock=None):
        '''
        Parse the data in the buffer

        If socket is provided, support defragmentation
        '''
        with self.lock:
            data.seek(0)
            offset = 0
            result = []

            if sock in self.defragmentation:
                save = self.defragmentation[sock]
                save.write(data.read())
                save.length += data.length
                # discard save
                data = save
                del self.defragmentation[sock]
                data.seek(0)

            while offset < data.length:
                # pick type and length
                (length, msg_type) = struct.unpack('IH', data.read(6))
                data.seek(offset)
                # if length + offset is greater than
                # remaining size, save the buffer for
                # defragmentation
                if (sock is not None) and (length + offset > data.length):
                    # create save buffer
                    self.defragmentation[sock] = save = io.BytesIO()
                    save.length = save.write(data.read())
                    # truncate data
                    data.truncate(offset)
                    break

                error = None
                if msg_type == NLMSG_ERROR:
                    data.seek(offset + 16)
                    code = abs(struct.unpack('i', data.read(4))[0])
                    if code > 0:
                        error = NetlinkError(code)
                    data.seek(offset)

                msg_class = self.msg_map.get(msg_type, nlmsg)
                msg = msg_class(data, debug=self.debug)
                try:
                    msg.decode()
                    msg['header']['error'] = error
                    # try to decode encapsulated error message
                    if error is not None:
                        enc_type = struct.unpack('H', msg.raw[24:26])[0]
                        enc_class = self.msg_map.get(enc_type, nlmsg)
                        enc = enc_class(msg.raw[20:])
                        enc.decode()
                        msg['header']['errmsg'] = enc
                except NetlinkHeaderDecodeError as e:
                    # in the case of header decoding error,
                    # create an empty message
                    msg = nlmsg()
                    msg['header']['error'] = e
                except NetlinkDecodeError as e:
                    msg['header']['error'] = e
                mtype = msg['header'].get('type', None)
                if mtype in (1, 2, 3, 4):
                    msg['event'] = mtypes.get(mtype, 'none')
                self.fix_message(msg)
                offset += msg.length
                result.append(msg)

            return result

    def fix_message(self, msg):
        pass

# 8<-----------------------------------------------------------
# Singleton, containing possible modifiers to the NetlinkSocket
# bind() call.
#
# Normally, you can open only one netlink connection for one
# process, but there is a hack. Current PID_MAX_LIMIT is 2^22,
# so we can use the rest to midify pid field.
#
# See also libnl library, lib/socket.c:generate_local_port()
sockets = AddrPool(minaddr=0x0,
                   maxaddr=0x3ff,
                   reverse=True)
# 8<-----------------------------------------------------------

class NetlinkSocket(socket.socket):
    '''
    Generic netlink socket
    '''

    def __init__(self, family=NETLINK_GENERIC, port=None, pid=None):
        socket.socket.__init__(self, socket.AF_NETLINK,
                               socket.SOCK_DGRAM, family)
        global sockets

        # 8<-----------------------------------------
        # PID init is here only for compatibility,
        # later it will be completely moved to bind()
        self.epid = None
        self.port = 0
        self.fixed = True
        if pid is None:
            self.pid = os.getpid() & 0x3fffff
            self.port = port
            self.fixed = self.port is not None
        elif pid == 0:
            self.pid = os.getpid()
        else:
            self.pid = pid
        # 8<-----------------------------------------
        self.groups = 0
        self.marshal = Marshal()

    def register_policy(self, policy, msg_class=None):
        '''
        Register netlink encoding/decoding policy. Can
        be specified in two ways:
        `nlsocket.register_policy(MSG_ID, msg_class)`
        to register one particular rule, or
        `nlsocket.register_policy({MSG_ID1: msg_class})`
        to register several rules at once.
        E.g.::

            policy = {RTM_NEWLINK: ifinfmsg,
                      RTM_DELLINK: ifinfmsg,
                      RTM_NEWADDR: ifaddrmsg,
                      RTM_DELADDR: ifaddrmsg}
            nlsocket.register_policy(policy)

        One can call `register_policy()` as many times,
        as one want to -- it will just extend the current
        policy scheme, not replace it.
        '''
        if isinstance(policy, int) and msg_class is not None:
            policy = {policy: msg_class}

        assert isinstance(policy, dict)
        for key in policy:
            self.marshal.msg_map[key] = policy[key]

        return self.marshal.msg_map

    def unregister_policy(self, policy):
        '''
        Unregister policy. Policy can be:

        * int -- then it will just remove one policy
        * list or tuple of ints -- remove all given
        * dict -- remove policies by keys from dict

        In the last case the routine will ignore dict values,
        it is implemented so just to make it compatible with
        `get_policy_map()` return value.
        '''
        if isinstance(policy, int):
            policy = [policy]
        elif isinstance(policy, dict):
            policy = list(policy)

        assert isinstance(policy, (tuple, list, set))

        for key in policy:
            del self.marshal.msg_map[key]

        return self.marshal.msg_map

    def get_policy_map(self, policy=None):
        '''
        Return policy for a given message type or for all
        message types. Policy parameter can be either int,
        or a list of ints. Always return dictionary.
        '''
        if policy is None:
            return self.marshal.msg_map

        if isinstance(policy, int):
            policy = [policy]

        assert isinstance(policy, (list, tuple, set))

        ret = {}
        for key in policy:
            ret[key] = self.marshal.msg_map[key]

        return ret

    def bind(self, groups=0, pid=None):
        '''
        Bind the socket to given multicast groups, using
        given pid.

        * If pid is None, use automatic port allocation
        * If pid == 0, use process' pid
        * If pid == <int>, use the value instead of pid
        '''
        if pid is not None:
            self.port = 0
            self.fixed = True
            self.pid = pid or os.getpid()

        self.groups = groups
        # if we have pre-defined port, use it strictly
        if self.fixed:
            self.epid = self.pid + (self.port << 22)
            socket.socket.bind(self, (self.epid, self.groups))
            return

        # if we have no pre-defined port, scan all the
        # range till the first available port
        for i in range(1024):
            try:
                self.port = sockets.alloc()
                self.epid = self.pid + (self.port << 22)
                socket.socket.bind(self, (self.epid, self.groups))
                # if we're here, bind() done successfully, just exit
                return
            except socket.error as e:
                # pass occupied sockets, raise other exceptions
                if e.errno != 98:
                    raise
        else:
            # raise "address in use" -- to be compatible
            raise socket.error(98, 'Address already in use')

    def get(self, bufsize=16384):
        '''
        Get parsed messages list.
        '''
        data = io.BytesIO()
        data.length = data.write(self.recv(bufsize))
        return self.marshal.parse(data, self)

    def close(self):
        '''
        Correctly close the socket and free all resources.
        '''
        global sockets
        if self.epid is not None:
            assert self.port is not None
            if not self.fixed:
                sockets.free(self.port)
            self.epid = None
        socket.socket.close(self)
<|endoftext|>