Skip to content

Commit

Permalink
Add handling for docs to plugin manager
Browse files Browse the repository at this point in the history
  • Loading branch information
Yngve S. Kristiansen authored and yngve-sk committed Jun 20, 2024
1 parent f123eac commit 4f363cc
Show file tree
Hide file tree
Showing 3 changed files with 84 additions and 18 deletions.
40 changes: 34 additions & 6 deletions src/ert/config/forward_model_step.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
import logging
from abc import abstractmethod
from dataclasses import dataclass, field
from typing import (
Dict,
List,
Literal,
Optional,
TypedDict,
Union,
)

from typing_extensions import NotRequired, Unpack
Expand Down Expand Up @@ -76,9 +79,27 @@ class ForwardModelStepOptions(TypedDict, total=False):
target_file: NotRequired[str]
error_file: NotRequired[str]
max_running_minutes: NotRequired[int]
environment: NotRequired[Dict[str, str]]
exec_env: NotRequired[Dict[str, str]]
default_mapping: NotRequired[Dict[str, str]]
environment: NotRequired[Dict[str, Union[str, int]]]
exec_env: NotRequired[Dict[str, Union[str, int]]]
default_mapping: NotRequired[Dict[str, Union[str, int]]]


@dataclass
class ForwardModelStepDocumentation:
config_file: Optional[str] = field(default=None)
source_package: str = field(default="ert")
source_function_name: str = field(default="ert")
description: str = field(default="No description")
examples: str = field(default="No examples")
category: Union[
Literal[
"utility.file_system",
"simulators.reservoir",
"modelling.reservoir",
"utility.templating",
],
str,
] = field(default="Uncategorized")


@dataclass
Expand Down Expand Up @@ -135,9 +156,9 @@ class ForwardModelStep:
arglist: List[str] = field(default_factory=list)
required_keywords: List[str] = field(default_factory=list)
arg_types: List[SchemaItemType] = field(default_factory=list)
environment: Dict[str, str] = field(default_factory=dict)
exec_env: Dict[str, str] = field(default_factory=dict)
default_mapping: Dict[str, str] = field(default_factory=dict)
environment: Dict[str, Union[int, str]] = field(default_factory=dict)
exec_env: Dict[str, Union[int, str]] = field(default_factory=dict)
default_mapping: Dict[str, Union[int, str]] = field(default_factory=dict)
private_args: SubstitutionList = field(default_factory=SubstitutionList)
help_text: str = ""

Expand Down Expand Up @@ -229,3 +250,10 @@ def __init__(
private_args=SubstitutionList(),
help_text="",
)

@staticmethod
@abstractmethod
def documentation() -> Optional[ForwardModelStepDocumentation]:
"""
Returns the documentation for the plugin forward model
"""
46 changes: 35 additions & 11 deletions src/ert/shared/_doc_utils/ert_jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,15 @@

from argparse import ArgumentParser
from collections import defaultdict
from typing import Any, Callable, Dict, List, Optional
from typing import Any, Callable, Dict, List, Optional, Union

import docutils.statemachine
from docutils import nodes
from docutils.statemachine import StringList
from sphinx.util.docutils import SphinxDirective
from sphinx.util.nodes import nested_parse_with_titles

from ert.config.forward_model_step import ForwardModelStepDocumentation
from ert.shared.plugins import ErtPluginManager, JobDoc


Expand All @@ -20,7 +21,7 @@ def __init__(
category: str,
job_source: str,
description: str,
job_config_file: str,
job_config_file: Optional[str],
parser: Optional[Callable[[], ArgumentParser]],
examples: Optional[str] = "",
) -> None:
Expand All @@ -37,10 +38,12 @@ def _create_job_config_section(self) -> nodes.section:
section_id=self.name + "-job-config", title="Job configuration used by ERT"
)

with open(self.job_config_file, encoding="utf-8") as fh:
job_config_text = fh.read()
job_config_text_node = nodes.literal_block(text=job_config_text)
config_section_node.append(job_config_text_node)
if self.job_config_file:
with open(self.job_config_file, encoding="utf-8") as fh:
job_config_text = fh.read()
job_config_text_node = nodes.literal_block(text=job_config_text)
config_section_node.append(job_config_text_node)

return config_section_node

def _create_job_details(self) -> nodes.definition_list:
Expand Down Expand Up @@ -99,8 +102,9 @@ def create_node(self, state: Any) -> nodes.section:
node.append(parser_section_node)

# Add forward model config file
config_section_node = self._create_job_config_section()
node.append(config_section_node)
if self.job_config_file:
config_section_node = self._create_job_config_section()
node.append(config_section_node)

return node

Expand All @@ -122,7 +126,9 @@ class _ErtDocumentation(SphinxDirective):

@staticmethod
def _divide_into_categories(
jobs: Dict[str, JobDoc],
jobs: Union[
Dict[str, JobDoc], Dict[str, Union[ForwardModelStepDocumentation, JobDoc]]
],
) -> Dict[str, Dict[str, List[_ForwardModelDocumentation]]]:
categories: Dict[str, Dict[str, List[_ForwardModelDocumentation]]] = (
defaultdict(lambda: defaultdict(list))
Expand All @@ -136,6 +142,16 @@ def _divide_into_categories(
if job_name.islower():
continue

if isinstance(docs, ForwardModelStepDocumentation):
docs = {
"description": docs.description,
"examples": docs.examples,
"config_file": docs.config_file,
"parser": None,
"source_package": docs.source_package,
"category": docs.category,
}

category = docs.get(
"category",
_ErtDocumentation._CATEGORY_DEFAULT,
Expand Down Expand Up @@ -183,7 +199,12 @@ def _create_forward_model_section_node(
return node

def _generate_job_documentation(
self, jobs: Dict[str, JobDoc], section_id: str, title: str
self,
jobs: Union[
Dict[str, JobDoc], Dict[str, Union[ForwardModelStepDocumentation, JobDoc]]
],
section_id: str,
title: str,
) -> List[nodes.section]:
job_categories = _ErtDocumentation._divide_into_categories(jobs)

Expand Down Expand Up @@ -242,7 +263,10 @@ def _create_section_with_title(section_id: str, title: str) -> nodes.section:

class ErtForwardModelDocumentation(_ErtDocumentation):
pm = ErtPluginManager()
_JOBS = pm.get_documentation_for_jobs()
_JOBS = {
**pm.get_documentation_for_jobs(),
**pm.get_documentation_for_forward_model_steps(),
}
_TITLE = "Pre-configured forward models"
_SECTION_ID = "ert-forward-models"

Expand Down
16 changes: 15 additions & 1 deletion src/ert/shared/plugins/plugin_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import pluggy

from ert.config.forward_model_step import (
ForwardModelStepDocumentation,
ForwardModelStepPlugin,
)
from ert.shared.plugins.workflow_config import WorkflowConfigs
Expand All @@ -54,7 +55,7 @@
class JobDoc(TypedDict):
description: str
examples: Optional[str]
config_file: str
config_file: Optional[str]
parser: Optional[Callable[[], ArgumentParser]]
source_package: str
category: str
Expand Down Expand Up @@ -280,6 +281,19 @@ def get_documentation_for_jobs(self) -> Dict[str, Any]:
)
return job_docs

def get_documentation_for_forward_model_steps(
self,
) -> Dict[str, ForwardModelStepDocumentation]:
return {
# Implementations of plugin fm step take no __init__ args
# (name, command)
# but mypy expects the subclasses to take in same arguments upon
# initializations
fm_step().name: fm_step.documentation() # type: ignore
for fm_step in self.forward_model_steps
if fm_step.documentation() is not None
}

def get_documentation_for_workflows(self) -> Dict[str, JobDoc]:
workflow_config = self.get_ertscript_workflows()

Expand Down

0 comments on commit 4f363cc

Please sign in to comment.