c11890829fe9fbf34031667537df7c0f313042fb
[idea/community.git] / python / src / com / jetbrains / python / codeInsight / postfix / PyBinaryConditionSurrounder.java
1 /*
2  * Copyright 2000-2016 JetBrains s.r.o.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package com.jetbrains.python.codeInsight.postfix;
17
18 import com.jetbrains.python.psi.*;
19 import org.jetbrains.annotations.Nullable;
20
21 public class PyBinaryConditionSurrounder extends PyExpressionAsConditionSurrounder {
22
23   private String myTextToGenerate;
24   private String myTemplateDescription;
25
26   public PyBinaryConditionSurrounder(String textToGenerate, String templateDescription) {
27     myTextToGenerate = textToGenerate;
28     myTemplateDescription = templateDescription;
29   }
30
31   @Override
32   protected String getTextToGenerate() {
33     return myTextToGenerate;
34   }
35
36   @Nullable
37   @Override
38   protected PyExpression getCondition(PyStatement statement) {
39     if (statement instanceof PyIfStatement) {
40       PyIfStatement ifStatement = (PyIfStatement)statement;
41       PyExpression condition = ifStatement.getIfPart().getCondition();
42       if (condition == null) {
43         return null;
44       }
45       return ((PyBinaryExpression)condition).getLeftExpression();
46     }
47     return null;
48   }
49
50   @Nullable
51   @Override
52   protected PyStatementListContainer getStatementListContainer(PyStatement statement) {
53     if (statement instanceof PyIfStatement) {
54       return ((PyIfStatement)statement).getIfPart();
55     }
56     return null;
57   }
58
59   @Override
60   public String getTemplateDescription() {
61     return myTemplateDescription;
62   }
63 }