From d2f548faf92cc31a796568a1bb1789dbab06fcd1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= Date: Fri, 12 May 2017 13:35:33 +0200 Subject: [PATCH] Proper type annotations for PillarServer.db() --- pillar/__init__.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/pillar/__init__.py b/pillar/__init__.py index 83d265e2..abe897da 100644 --- a/pillar/__init__.py +++ b/pillar/__init__.py @@ -7,6 +7,7 @@ import logging import logging.config import subprocess import tempfile +import typing import os import os.path @@ -15,6 +16,8 @@ from eve import Eve import flask from flask import render_template, request from flask.templating import TemplateNotFound +import pymongo.collection +import pymongo.database from pillar.api import custom_field_validation from pillar.api.utils import authentication @@ -551,11 +554,9 @@ class PillarServer(Eve): pprint(links) - def db(self, collection_name: str=None): - """Returns the MongoDB database. - - :rtype: flask_pymongo.PyMongo - """ + def db(self, collection_name: str=None) \ + -> typing.Union[pymongo.collection.Collection, pymongo.database.Database]: + """Returns the MongoDB database, or the collection (if given)""" if collection_name: return self.data.driver.db[collection_name]