datasources.py

# Based on: https://github.com/allenai/s2-folks/tree/main/examples/python
#
#

import os
import re

FORMAT_TEXT = "Text"
FORMAT_DICT = "Dict"
FORMAT_JSON = "JSON"

FORMATS = (FORMAT_TEXT, FORMAT_DICT)

DS_WORMNEUROATLAS = "WormNeuroAtlas"

import wormneuroatlas as wa
import pprint

pp = pprint.PrettyPrinter(depth=6)


def query_data_source(text, source, format):

    if source == DS_WORMNEUROATLAS:

        atlas = wa.NeuroAtlas()
        ds, ss = atlas.everything_about(
            text, return_values=True, return_text=True, print_text=False
        )

        if format == FORMAT_TEXT:
            info = """Extracting info on neuron %s from %s...
""" % (
                text,
                source,
            )
            info += """```%s```""" % ss
        elif format == FORMAT_DICT:
            info = """```
%s```""" % pp.pformat(
                ds
            )
        elif format == FORMAT_JSON:
            import json

            jj = json.dumps(ds)
            info = """```
%s```""" % pp.pformat(
                jj
            )

    return info