From d7f73eef6850018a00729ea0be3025ed347c4b92 Mon Sep 17 00:00:00 2001 From: James Mathews Date: Thu, 27 Jul 2023 18:10:35 -0400 Subject: [PATCH] Ensure cursor still around for repeat use of feature matrix extractor object. --- .../db/feature_matrix_extractor.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/spatialprofilingtoolbox/db/feature_matrix_extractor.py b/spatialprofilingtoolbox/db/feature_matrix_extractor.py index 3f005221f..0a096021f 100644 --- a/spatialprofilingtoolbox/db/feature_matrix_extractor.py +++ b/spatialprofilingtoolbox/db/feature_matrix_extractor.py @@ -2,7 +2,6 @@ Convenience provision of a feature matrix for each study, the data retrieved from the SPT database. """ import sys -from typing import cast import pandas as pd from psycopg2.extensions import cursor as Psycopg2Cursor @@ -22,7 +21,7 @@ class FeatureMatrixExtractor: Pull from the database and create convenience bundle of feature matrices and metadata. """ - cursor: Psycopg2Cursor + cursor: Psycopg2Cursor | None database_config_file: str | None def __init__(self, @@ -36,7 +35,7 @@ def __init__(self, message = 'A cursor and database configuration file were both specified. Using the '\ 'cursor.' logger.warning(message) - self.cursor = cast(Psycopg2Cursor, cursor) + self.cursor = cursor self.database_config_file = database_config_file def extract(self, @@ -44,14 +43,16 @@ def extract(self, study: str | None=None, continuous_also: bool=False, ): + extraction = None if self.cursor is not None: - return self._extract(specimen=specimen, study=study, continuous_also=continuous_also) + extraction = self._extract(specimen=specimen, study=study, continuous_also=continuous_also) if self.database_config_file is not None: with DatabaseConnectionMaker(self.database_config_file) as dcm: with dcm.get_connection().cursor() as cursor: self.cursor = cursor - return self._extract(specimen=specimen, study=study, continuous_also=continuous_also) - return None + extraction = self._extract(specimen=specimen, study=study, continuous_also=continuous_also) + self.cursor = None + return extraction def _extract(self, specimen: str | None=None,