Merge pull request #340 from equinor/add-stubs

Add stubs

Author: EvenSol

.github/workflows/generate-stubs.yml

@@ -0,0 +1,45 @@
+name: Generate Java Stubs
+
+on:
+  # Run when JAR files are updated
+  push:
+    paths:
+      - 'src/neqsim/lib/**/*.jar'
+  # Run on release
+  release:
+    types: [published]
+  # Allow manual trigger
+  workflow_dispatch:
+
+jobs:
+  generate-stubs:
+    runs-on: ubuntu-latest
+    
+    steps:
+      - uses: actions/checkout@v4
+      
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      
+      - name: Set up Java
+        uses: actions/setup-java@v4
+        with:
+          distribution: 'temurin'
+          java-version: '17'
+      
+      - name: Install dependencies
+        run: |
+          pip install poetry
+          poetry install --with dev
+      
+      - name: Generate stubs
+        run: |
+          poetry run python scripts/generate_stubs.py
+      
+      - name: Commit updated stubs
+        uses: stefanzweifel/git-auto-commit-action@v5
+        with:
+          commit_message: "chore: regenerate Java stubs"
+          file_pattern: "src/jneqsim-stubs/**"

pyproject.toml

@@ -27,6 +27,7 @@ tabulate = { version = "^0.9.0", optional = true }
 black = ">=23.12,<25.0"
 pytest = "^7.4.3"
 pre-commit = "^3.5.0" # Higher versions require python 3.9+
+stubgenj = "^0.2.12"  # Generate type stubs for Java classes
 
 [tool.poetry.extras]
 interactive = ["matplotlib", "jupyter", "tabulate"]

scripts/generate_stubs.py

@@ -0,0 +1,117 @@
+"""
+Script to generate Python type stubs for neqsim Java classes using stubgenj.
+
+This enables IDE autocompletion and type checking for the neqsim Java library
+accessed via JPype.
+
+The Java package 'neqsim' is exposed as 'jneqsim' in Python to avoid naming
+conflicts with the Python 'neqsim' package. The stubs are generated accordingly.
+
+Usage:
+    python scripts/generate_stubs.py
+
+The stubs will be generated in the src/jneqsim-stubs directory.
+"""
+
+import re
+import shutil
+import sys
+from pathlib import Path
+
+# Add src to path
+src_path = Path(__file__).parent.parent / "src"
+sys.path.insert(0, str(src_path))
+
+
+def rename_package_in_stubs(stubs_dir: Path, old_name: str, new_name: str):
+    """
+    Rename all references from old_name to new_name in stub files.
+    This handles the neqsim -> jneqsim renaming to avoid conflicts
+    with the Python neqsim package.
+    """
+    for pyi_file in stubs_dir.rglob("*.pyi"):
+        content = pyi_file.read_text(encoding="utf-8")
+
+        # Replace import statements and type references
+        # Match 'neqsim.' but not 'jneqsim.' (negative lookbehind)
+        new_content = re.sub(
+            rf"(?<!j){old_name}\.", f"{new_name}.", content
+        )
+
+        if new_content != content:
+            pyi_file.write_text(new_content, encoding="utf-8")
+
+
+def generate_stubs():
+    """Generate type stubs for neqsim Java classes."""
+    import jpype
+    import jpype.imports  # Enable Java imports
+    import stubgenj
+
+    # Start JVM if not already started
+    if not jpype.isJVMStarted():
+        # Import neqsim to start JVM with correct classpath
+        import neqsim  # noqa: F401
+
+    # Import the neqsim Java package to get JPackage reference
+    from neqsim.neqsimpython import jneqsim
+
+    # Temporary output directory (stubgenj will create 'neqsim-stubs')
+    temp_output_dir = src_path / "temp-stubs"
+    if temp_output_dir.exists():
+        shutil.rmtree(temp_output_dir)
+    temp_output_dir.mkdir(exist_ok=True)
+
+    # Final output directory
+    final_output_dir = src_path / "jneqsim-stubs"
+
+    print("Generating stubs...")
+
+    # Generate stubs for the neqsim package (pass JPackage objects)
+    stubgenj.generateJavaStubs(
+        parentPackages=[jneqsim],  # The neqsim JPackage
+        useStubsSuffix=True,  # Creates neqsim-stubs folder structure
+        outputDir=str(temp_output_dir),
+        jpypeJPackageStubs=True,  # Include jpype stubs
+        includeJavadoc=True,  # Include javadoc in stubs
+    )
+
+    # Rename neqsim -> jneqsim in all stub files to avoid conflict
+    # with Python neqsim package
+    neqsim_stubs = temp_output_dir / "neqsim-stubs"
+    if neqsim_stubs.exists():
+        print("Renaming 'neqsim' -> 'jneqsim' in stubs to avoid naming conflict...")
+        rename_package_in_stubs(temp_output_dir, "neqsim", "jneqsim")
+
+        # Clean up existing output
+        if final_output_dir.exists():
+            shutil.rmtree(final_output_dir)
+        final_output_dir.mkdir(exist_ok=True)
+
+        # Move jpype-stubs as-is (it's at temp_output_dir/jpype-stubs)
+        jpype_stubs = temp_output_dir / "jpype-stubs"
+        if jpype_stubs.exists():
+            shutil.move(str(jpype_stubs), str(final_output_dir / "jpype-stubs"))
+
+        # Rename folder neqsim-stubs -> jneqsim-stubs
+        target = final_output_dir / "jneqsim-stubs"
+        shutil.move(str(neqsim_stubs), str(target))
+
+        # Clean up temp directory
+        shutil.rmtree(temp_output_dir)
+
+    print(f"Stubs generated successfully in {final_output_dir}")
+    print("\n" + "=" * 60)
+    print("USAGE INSTRUCTIONS")
+    print("=" * 60)
+    print("\nThe Java 'neqsim' package stubs are available as 'jneqsim'")
+    print("to avoid conflicts with the Python 'neqsim' package.")
+    print("\nFor VS Code with Pylance, add to settings.json:")
+    print('  "python.analysis.extraPaths": ["src/jneqsim-stubs"]')
+    print("\nFor mypy, add to pyproject.toml:")
+    print('  [tool.mypy]')
+    print('  mypy_path = "src/jneqsim-stubs"')
+
+
+if __name__ == "__main__":
+    generate_stubs()

src/jneqsim-stubs/jneqsim-stubs/__init__.pyi

@@ -0,0 +1,41 @@
+
+import sys
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
+import jneqsim.api
+import jneqsim.blackoil
+import jneqsim.chemicalreactions
+import jneqsim.datapresentation
+import jneqsim.fluidmechanics
+import jneqsim.mathlib
+import jneqsim.physicalproperties
+import jneqsim.process
+import jneqsim.pvtsimulation
+import jneqsim.standards
+import jneqsim.statistics
+import jneqsim.thermo
+import jneqsim.thermodynamicoperations
+import jneqsim.util
+import typing
+
+
+class __module_protocol__(Protocol):
+    # A module protocol which reflects the result of ``jp.JPackage("neqsim")``.
+
+    api: jneqsim.api.__module_protocol__
+    blackoil: jneqsim.blackoil.__module_protocol__
+    chemicalreactions: jneqsim.chemicalreactions.__module_protocol__
+    datapresentation: jneqsim.datapresentation.__module_protocol__
+    fluidmechanics: jneqsim.fluidmechanics.__module_protocol__
+    mathlib: jneqsim.mathlib.__module_protocol__
+    physicalproperties: jneqsim.physicalproperties.__module_protocol__
+    process: jneqsim.process.__module_protocol__
+    pvtsimulation: jneqsim.pvtsimulation.__module_protocol__
+    standards: jneqsim.standards.__module_protocol__
+    statistics: jneqsim.statistics.__module_protocol__
+    thermo: jneqsim.thermo.__module_protocol__
+    thermodynamicoperations: jneqsim.thermodynamicoperations.__module_protocol__
+    util: jneqsim.util.__module_protocol__

src/jneqsim-stubs/jneqsim-stubs/api/__init__.pyi

@@ -0,0 +1,15 @@
+
+import sys
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
+import jneqsim.api.ioc
+import typing
+
+
+class __module_protocol__(Protocol):
+    # A module protocol which reflects the result of ``jp.JPackage("jneqsim.api")``.
+
+    ioc: jneqsim.api.ioc.__module_protocol__

src/jneqsim-stubs/jneqsim-stubs/api/ioc/__init__.pyi

@@ -0,0 +1,25 @@
+
+import sys
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
+import java.lang
+import jpype
+import typing
+
+
+
+class CalculationResult:
+    fluidProperties: typing.MutableSequence[typing.MutableSequence[float]] = ...
+    calculationError: typing.MutableSequence[java.lang.String] = ...
+    def __init__(self, doubleArray: typing.Union[typing.List[typing.MutableSequence[float]], jpype.JArray], stringArray: typing.Union[typing.List[java.lang.String], jpype.JArray]): ...
+    def equals(self, object: typing.Any) -> bool: ...
+    def hashCode(self) -> int: ...
+
+
+class __module_protocol__(Protocol):
+    # A module protocol which reflects the result of ``jp.JPackage("jneqsim.api.ioc")``.
+
+    CalculationResult: typing.Type[CalculationResult]

src/jneqsim-stubs/jneqsim-stubs/blackoil/__init__.pyi

@@ -0,0 +1,113 @@
+
+import sys
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
+import java.util
+import jpype
+import jneqsim.blackoil.io
+import jneqsim.thermo.system
+import typing
+
+
+
+class BlackOilConverter:
+    def __init__(self): ...
+    @staticmethod
+    def convert(systemInterface: jneqsim.thermo.system.SystemInterface, double: float, doubleArray: typing.Union[typing.List[float], jpype.JArray], double3: float, double4: float) -> 'BlackOilConverter.Result': ...
+    class Result:
+        pvt: 'BlackOilPVTTable' = ...
+        blackOilSystem: 'SystemBlackOil' = ...
+        rho_o_sc: float = ...
+        rho_g_sc: float = ...
+        rho_w_sc: float = ...
+        bubblePoint: float = ...
+        def __init__(self): ...
+
+class BlackOilFlash:
+    def __init__(self, blackOilPVTTable: 'BlackOilPVTTable', double: float, double2: float, double3: float): ...
+    def flash(self, double: float, double2: float, double3: float, double4: float, double5: float) -> 'BlackOilFlashResult': ...
+
+class BlackOilFlashResult:
+    O_std: float = ...
+    Gf_std: float = ...
+    W_std: float = ...
+    V_o: float = ...
+    V_g: float = ...
+    V_w: float = ...
+    rho_o: float = ...
+    rho_g: float = ...
+    rho_w: float = ...
+    mu_o: float = ...
+    mu_g: float = ...
+    mu_w: float = ...
+    Rs: float = ...
+    Rv: float = ...
+    Bo: float = ...
+    Bg: float = ...
+    Bw: float = ...
+    def __init__(self): ...
+
+class BlackOilPVTTable:
+    def __init__(self, list: java.util.List['BlackOilPVTTable.Record'], double: float): ...
+    def Bg(self, double: float) -> float: ...
+    def Bo(self, double: float) -> float: ...
+    def Bw(self, double: float) -> float: ...
+    def Rs(self, double: float) -> float: ...
+    def RsEffective(self, double: float) -> float: ...
+    def Rv(self, double: float) -> float: ...
+    def getBubblePointP(self) -> float: ...
+    def mu_g(self, double: float) -> float: ...
+    def mu_o(self, double: float) -> float: ...
+    def mu_w(self, double: float) -> float: ...
+    class Record:
+        p: float = ...
+        Rs: float = ...
+        Bo: float = ...
+        mu_o: float = ...
+        Bg: float = ...
+        mu_g: float = ...
+        Rv: float = ...
+        Bw: float = ...
+        mu_w: float = ...
+        def __init__(self, double: float, double2: float, double3: float, double4: float, double5: float, double6: float, double7: float, double8: float, double9: float): ...
+
+class SystemBlackOil:
+    def __init__(self, blackOilPVTTable: BlackOilPVTTable, double: float, double2: float, double3: float): ...
+    def copyShallow(self) -> 'SystemBlackOil': ...
+    def flash(self) -> BlackOilFlashResult: ...
+    def getBg(self) -> float: ...
+    def getBo(self) -> float: ...
+    def getBw(self) -> float: ...
+    def getGasDensity(self) -> float: ...
+    def getGasReservoirVolume(self) -> float: ...
+    def getGasStdTotal(self) -> float: ...
+    def getGasViscosity(self) -> float: ...
+    def getOilDensity(self) -> float: ...
+    def getOilReservoirVolume(self) -> float: ...
+    def getOilStdTotal(self) -> float: ...
+    def getOilViscosity(self) -> float: ...
+    def getPressure(self) -> float: ...
+    def getRs(self) -> float: ...
+    def getRv(self) -> float: ...
+    def getTemperature(self) -> float: ...
+    def getWaterDensity(self) -> float: ...
+    def getWaterReservoirVolume(self) -> float: ...
+    def getWaterStd(self) -> float: ...
+    def getWaterViscosity(self) -> float: ...
+    def setPressure(self, double: float) -> None: ...
+    def setStdTotals(self, double: float, double2: float, double3: float) -> None: ...
+    def setTemperature(self, double: float) -> None: ...
+
+
+class __module_protocol__(Protocol):
+    # A module protocol which reflects the result of ``jp.JPackage("jneqsim.blackoil")``.
+
+    BlackOilConverter: typing.Type[BlackOilConverter]
+    BlackOilFlash: typing.Type[BlackOilFlash]
+    BlackOilFlashResult: typing.Type[BlackOilFlashResult]
+    BlackOilPVTTable: typing.Type[BlackOilPVTTable]
+    SystemBlackOil: typing.Type[SystemBlackOil]
+    io: jneqsim.blackoil.io.__module_protocol__