diff --git a/modelscan/cli.py b/modelscan/cli.py index a10fcb9..405d936 100644 --- a/modelscan/cli.py +++ b/modelscan/cli.py @@ -178,17 +178,21 @@ def create_settings(force: bool, location: Optional[str]) -> None: if force: with open(settings_path, mode="w", encoding="utf-8") as settings_file: settings_file.write(SettingsUtils.get_default_settings_as_toml()) + settings_file.close() else: logger.warning( "%s file already exists. Please use `--force` flag if you intend to overwrite it.", settings_path, ) + except FileNotFoundError: with open(settings_path, mode="w", encoding="utf-8") as settings_file: settings_file.write(SettingsUtils.get_default_settings_as_toml()) + settings_file.close() def main() -> None: + result = 0 try: result = cli.main(standalone_mode=False) diff --git a/modelscan/issues.py b/modelscan/issues.py index 16bfb51..b254084 100644 --- a/modelscan/issues.py +++ b/modelscan/issues.py @@ -132,6 +132,7 @@ def __init__( source: Union[Path, str], scanner: str = "", ) -> None: + super().__init__(scanner) self.module = module self.operator = operator self.source = source diff --git a/modelscan/model.py b/modelscan/model.py index 22b94ca..e6a5f4b 100644 --- a/modelscan/model.py +++ b/modelscan/model.py @@ -9,13 +9,13 @@ class ModelDataEmpty(ValueError): class Model: _source: Path _stream: Optional[IO[bytes]] - _source_file_used: bool + _should_close_stream: bool # Flag to control closing of file _context: Dict[str, Any] def __init__(self, source: Union[str, Path], stream: Optional[IO[bytes]] = None): self._source = Path(source) self._stream = stream - self._source_file_used = False + self._should_close_stream = stream is None # Only close if opened self._context = {"formats": []} def set_context(self, key: str, value: Any) -> None: @@ -29,14 +29,16 @@ def open(self) -> "Model": return self self._stream = open(self._source, "rb") - self._source_file_used = True + self._should_close_stream = True return self def close(self) -> None: # Only close the stream if we opened a file (not for IO[bytes] objects passed in) - if self._stream and self._source_file_used: + if self._stream and self._should_close_stream: self._stream.close() + self._stream = None # Avoid double-closing + self._should_close_stream = False # Reset the flag def __enter__(self) -> "Model": return self.open() diff --git a/modelscan/skip.py b/modelscan/skip.py index 2f83b75..73d0885 100644 --- a/modelscan/skip.py +++ b/modelscan/skip.py @@ -1,5 +1,4 @@ import logging -from enum import Enum from modelscan.settings import Property diff --git a/modelscan/tools/picklescanner.py b/modelscan/tools/picklescanner.py index 95bddd6..44c4e2a 100644 --- a/modelscan/tools/picklescanner.py +++ b/modelscan/tools/picklescanner.py @@ -5,7 +5,7 @@ import numpy as np -from modelscan.error import ModelScanError, PickleGenopsError +from modelscan.error import PickleGenopsError from modelscan.skip import ModelScanSkipped, SkipCategories from modelscan.issues import Issue, IssueCode, IssueSeverity, OperatorIssueDetails from modelscan.scanners.scan import ScanResults