Add useful type hints for possibly a better experience for API users

[sinaatalay]

Hi, this PR is like a question from me. The only change in it is adding a return type to the __getitem__ method of the desc.equilibrium.EquilibriaFamily class.

I started learning DESC, and I thought type hints may make the API a little easier to use.

For example, I wanted to get an equilibrium from the instance of a EquilibriaFamily class by using the indexing operator like below:

equilibria_family = desc.continuation.solve_continuation_automatic(equilibrium)
an_equilibrium = family_of_equilibria[-1]

But, since the __getitem__ method of the EquilibriaFamily class doesn't have a return type, my IDE doesn't know that the variable an_equilibrium is an instance of Equilibrium. So, I have to tell it myself to see the available methods of the Equilibrium class, such as the save method.

Would you be interested if I added useful types like this as I come across them to improve the API experience?

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.62%. Comparing base (2741269) to head (6ac29dc).
Report is 855 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1344      +/-   ##
==========================================
+ Coverage   95.57%   95.62%   +0.05%     
==========================================
  Files          96       98       +2     
  Lines       24405    25420    +1015     
==========================================
+ Hits        23325    24309     +984     
- Misses       1080     1111      +31     

Files with missing lines	Coverage Δ
desc/continuation.py	93.26% <ø> (ø)
desc/equilibrium/equilibrium.py	96.07% <ø> (ø)

... and 40 files with indirect coverage changes

[YigitElma]

Benchmarks fail due to this PR coming from a fork, I recommend you to use the original repo since you can 😄 . But #1345 should solve the issue for other contributors.

[dpanici]

I think the issue here is that we would need a lot of local imports to do this across all our fxns, and some of those may lead to circular import errors because we have a lot of interdependencies in our files

[sinaatalay]

I think the issue here is that we would need a lot of local imports to do this across all our fxns, and some of those may lead to circular import errors because we have a lot of interdependencies in our files

Oh, right, yes, circular imports will be very likely then.

If we decide to provide types in the future anyway, I am a volunteer.

I will close the PR in a couple of days to see if maybe someone else has a comment on this.

[dpanici]

can add "if type checking" block to top?

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     -2.19 +/- 4.93     | -1.37e-02 +/- 3.07e-02 |  6.09e-01 +/- 2.6e-02  |  6.23e-01 +/- 1.6e-02  |
 test_build_transform_fft_highres        |     -1.69 +/- 3.36     | -1.66e-02 +/- 3.30e-02 |  9.65e-01 +/- 2.3e-02  |  9.81e-01 +/- 2.3e-02  |
 test_equilibrium_init_lowres            |     -1.56 +/- 3.40     | -6.18e-02 +/- 1.35e-01 |  3.89e+00 +/- 5.1e-02  |  3.96e+00 +/- 1.2e-01  |
 test_objective_compile_atf              |     +0.43 +/- 4.08     | +3.35e-02 +/- 3.20e-01 |  7.87e+00 +/- 2.6e-01  |  7.84e+00 +/- 1.9e-01  |
 test_objective_compute_atf              |     +8.83 +/- 3.06     | +9.22e-04 +/- 3.20e-04 |  1.14e-02 +/- 2.9e-04  |  1.04e-02 +/- 1.3e-04  |
 test_objective_jac_atf                  |     +3.05 +/- 2.27     | +5.78e-02 +/- 4.29e-02 |  1.95e+00 +/- 3.1e-02  |  1.89e+00 +/- 2.9e-02  |
 test_perturb_1                          |     +4.22 +/- 2.56     | +5.77e-01 +/- 3.50e-01 |  1.42e+01 +/- 3.1e-01  |  1.37e+01 +/- 1.7e-01  |
 test_proximal_jac_atf                   |     +0.80 +/- 0.98     | +6.49e-02 +/- 7.91e-02 |  8.15e+00 +/- 4.2e-02  |  8.09e+00 +/- 6.7e-02  |
 test_proximal_freeb_compute             |     +1.52 +/- 1.22     | +2.96e-03 +/- 2.36e-03 |  1.97e-01 +/- 2.1e-03  |  1.94e-01 +/- 1.0e-03  |
 test_build_transform_fft_lowres         |     +9.28 +/- 4.75     | +5.02e-02 +/- 2.57e-02 |  5.91e-01 +/- 1.7e-02  |  5.41e-01 +/- 1.9e-02  |
 test_equilibrium_init_medres            |     +5.54 +/- 3.17     | +2.38e-01 +/- 1.36e-01 |  4.53e+00 +/- 1.1e-01  |  4.29e+00 +/- 7.8e-02  |
 test_equilibrium_init_highres           |     +2.66 +/- 1.88     | +1.48e-01 +/- 1.05e-01 |  5.71e+00 +/- 7.3e-02  |  5.56e+00 +/- 7.5e-02  |
 test_objective_compile_dshape_current   |     +0.95 +/- 2.02     | +3.71e-02 +/- 7.87e-02 |  3.93e+00 +/- 6.2e-02  |  3.89e+00 +/- 4.8e-02  |
 test_objective_compute_dshape_current   |     +0.27 +/- 2.31     | +9.74e-06 +/- 8.25e-05 |  3.58e-03 +/- 6.5e-05  |  3.57e-03 +/- 5.1e-05  |
 test_objective_jac_dshape_current       |     -7.15 +/- 5.88     | -3.11e-03 +/- 2.56e-03 |  4.04e-02 +/- 1.7e-03  |  4.35e-02 +/- 1.9e-03  |
 test_perturb_2                          |     -1.28 +/- 3.70     | -2.51e-01 +/- 7.26e-01 |  1.94e+01 +/- 3.3e-01  |  1.96e+01 +/- 6.5e-01  |
 test_proximal_freeb_jac                 |     +0.50 +/- 0.99     | +3.70e-02 +/- 7.38e-02 |  7.50e+00 +/- 5.5e-02  |  7.46e+00 +/- 4.9e-02  |
 test_solve_fixed_iter                   |     +0.41 +/- 58.25    | +2.06e-02 +/- 2.95e+00 |  5.08e+00 +/- 2.1e+00  |  5.06e+00 +/- 2.1e+00  |

[YigitElma]

@sinaatalay do you plan to work on this?

[sinaatalay]

Yes, I will start working on it soon.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[dpanici]

For both outputs and inputs

• Equilibrium
• ObjectiveFunctions
• Coils
• Plotting (can return fig, ax and also can return data)
• Grid
• Geometry (surface, grid)