From d4dc6a9e88c6019d3f456ec2f7e524886b73f476 Mon Sep 17 00:00:00 2001
From: Nicholas McKibben <nicholas.bgp@gmail.com>
Date: Fri, 24 Dec 2021 21:05:14 -0700
Subject: [PATCH] ENH: support writing MIP MPS files via integrality argument

---
 glpk/_fileio.py | 3 ++-
 glpk/_utils.py  | 8 ++++++--
 setup.py        | 4 ++--
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/glpk/_fileio.py b/glpk/_fileio.py
index e281c04..b37d9ff 100644
--- a/glpk/_fileio.py
+++ b/glpk/_fileio.py
@@ -144,13 +144,14 @@ def mpswrite(
         A_eq=None,
         b_eq=None,
         bounds=None,
+        integrality=None,
         sense=GLPK.GLP_MIN,
         filename='prob.mps',
         fmt=GLPK.GLP_MPS_FILE):
     '''Write an MPS file.'''
 
     filename = pathlib.Path(filename)
-    prob, _lp = _fill_prob(c, A_ub, b_ub, A_eq, b_eq, bounds, sense, filename.stem)
+    prob, _lp = _fill_prob(c, A_ub, b_ub, A_eq, b_eq, bounds, integrality, sense, filename.stem)
 
     # Get the library
     _lib = GLPK()._lib
diff --git a/glpk/_utils.py b/glpk/_utils.py
index 33cd312..7b99d49 100644
--- a/glpk/_utils.py
+++ b/glpk/_utils.py
@@ -6,7 +6,7 @@
 from scipy.sparse import coo_matrix
 from scipy.optimize._linprog_util import _LPProblem, _clean_inputs
 
-from ._glpk_defines import GLPK, glp_prob
+from ._glpk_defines import GLPK
 
 
 def _convert_bounds(processed_bounds):
@@ -30,10 +30,11 @@ def _convert_bounds(processed_bounds):
     return bounds
 
 
-def _fill_prob(c, A_ub, b_ub, A_eq, b_eq, bounds, sense, prob_name):
+def _fill_prob(c, A_ub, b_ub, A_eq, b_eq, bounds, integrality, sense, prob_name):
     '''Create and populate GLPK prob struct from linprog definition.'''
 
     # Housekeeping
+    # TODO: modify to use integrality after https://github.com/mckib2/scipy/pull/28 makes it into scipy master
     lp = _clean_inputs(_LPProblem(c, A_ub, b_ub, A_eq, b_eq, bounds, None))
     c, A_ub, b_ub, A_eq, b_eq, processed_bounds, _x0 = lp
 
@@ -68,6 +69,9 @@ def _fill_prob(c, A_ub, b_ub, A_eq, b_eq, bounds, sense, prob_name):
             _lib.glp_set_col_bnds(prob, ii + first_col, bnd[0], bnd[1], bnd[2])
         # else: default is GLP_FX with lb=0, ub=0
 
+        if integrality and integrality[ii] == 1:
+            _lib.glp_set_col_kind(prob, ii + first_col, GLPK.GLP_IV)
+
     # Need to load both matrices at the same time
     first_row = _lib.glp_add_rows(prob, A.shape[0])
 
diff --git a/setup.py b/setup.py
index 4e29509..7d09c64 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@ def scrape_makefile_list(filename, START_TOKEN, END_TOKEN):
 
 
 class build_ext(_build_ext):
-    '''Overide get_export_symbols to provide them for Windows DLL.'''
+    '''Override get_export_symbols to provide them for Windows DLL.'''
     def get_export_symbols(self, ext):
         '''Only for generating Windows DLL.'''
         def_file = GLPK_SRC_DIR / '../w64/glpk_4_65.def'
@@ -41,7 +41,7 @@ def get_export_symbols(self, ext):
 
 setup(
     name='scikit-glpk',
-    version='0.4.3',
+    version='0.4.4',
     author='Nicholas McKibben',
     author_email='nicholas.bgp@gmail.com',
     url='https://github.com/mckib2/scikit-glpk',