Custom functions and pipelines
To include this logic into a pipeline you have to create a custom transformer. You need to ask yourself: 


[INIT] Are there any parameters in my logic? 


The variable you want to impute and the category you want this imputation to be based on.

[FIT] What part of the logic is related to computing what the transformation will be?


When you compute the median() by groups and store the data somehow for later transform.

[TRANSFORM] What part of the logic transforms the data, given the parameters (in 1) and the setting that was made (in 2)?


When you get the parameters (access specific key in the dictionary) to retrieve what was the mean for that group, and then fill the missing value with this.



Here is an example : 

from sklearn.base import BaseEstimator, TransformerMixin

class CustomImputer(BaseEstimator, TransformerMixin) : 
     def __init__(self, variable, by) : 
          #self.something enables you to include the passed parameters
          #as object attributes and use it in other methods of the class
          self.variable = variable
          self.by = by
          return self

     def fit(self, X, y=None) : 
          self.map = X.groupby(self.by)[variable].mean()
          #self.map become an attribute that is, the map of values to
          #impute in function of index (corresponding table, like a dict)
          return self

     def transform(self, X, y=None) : 
          X[variable] = X[variable].fillna(value = X[by].map(self.map))
          #Change the variable column. If the value is missing, value should 
          #be replaced by the mapping of column "by" according to the map you
          #created in fit method (self.map)
          return X


Now, it can be included in any pipeline :

#Minimal example, you could include this imputer in columns transformer to 
#apply it multiple time
pipeline = Pipeline(steps = [('myImputer', CustomImputer('variabletofill',
                                                         'based_on_variable'),
                              ('model', LinearRegression())])

y_pred = pipeline.fit(X_train, y_train).predict(X_test)


As you can see, the mapping is computed only based on train data. Then it is re-used to impute the missing values. It is data-leakage-proof.
Here is a good article that explains how to create a custom transformer.

Hope this helps